Merge pull request 'docs(clawhdf5): document DType variants, fix unresolved doc links' (#17) from sdlc-docs/clawhdf5-types-20260514-165210 into main
This commit is contained in:
redclawsystems
@@ -0,0 +1,276 @@
|
||||
use clawhdf5::writer::FileBuilder;
|
||||
use clawhdf5_format::datatype::{CharacterSet, Datatype, StringPadding};
|
||||
use clawhdf5_format::type_builders::AttrValue;
|
||||
|
||||
use crate::sqlite_reader::SqliteData;
|
||||
|
||||
/// Options controlling HDF5 output.
|
||||
pub struct WriteOptions {
|
||||
pub agent_id: String,
|
||||
pub embedder: String,
|
||||
pub compression: bool,
|
||||
pub compression_level: u32,
|
||||
pub float16: bool,
|
||||
}
|
||||
|
||||
/// Write SQLite data to an HDF5 file.
|
||||
pub fn write_hdf5(
|
||||
path: &str,
|
||||
data: &SqliteData,
|
||||
opts: &WriteOptions,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let mut builder = FileBuilder::new();
|
||||
|
||||
// Root-level metadata attributes
|
||||
builder.set_attr("agent_id", AttrValue::String(opts.agent_id.clone()));
|
||||
builder.set_attr("embedder", AttrValue::String(opts.embedder.clone()));
|
||||
builder.set_attr("embedding_dim", AttrValue::I64(data.embedding_dim as i64));
|
||||
builder.set_attr("source", AttrValue::String("sqlite-migration".into()));
|
||||
builder.set_attr("version", AttrValue::I64(1));
|
||||
|
||||
write_chunks_group(&mut builder, data, opts);
|
||||
write_sessions_group(&mut builder, data);
|
||||
write_entities_group(&mut builder, data);
|
||||
write_relations_group(&mut builder, data);
|
||||
|
||||
builder.write(path)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Build a fixed-length string Datatype from the max byte length of the items.
|
||||
fn string_dtype(max_len: usize) -> Datatype {
|
||||
Datatype::String {
|
||||
size: max_len.max(1) as u32,
|
||||
padding: StringPadding::NullPad,
|
||||
charset: CharacterSet::Utf8,
|
||||
}
|
||||
}
|
||||
|
||||
/// Pack a slice of strings into null-padded raw bytes of uniform width.
|
||||
fn pack_strings(strings: &[String]) -> (Vec<u8>, usize) {
|
||||
let max_len = strings.iter().map(|s| s.len()).max().unwrap_or(0).max(1);
|
||||
let mut buf = vec![0u8; strings.len() * max_len];
|
||||
for (i, s) in strings.iter().enumerate() {
|
||||
let start = i * max_len;
|
||||
let bytes = s.as_bytes();
|
||||
let copy_len = bytes.len().min(max_len);
|
||||
buf[start..start + copy_len].copy_from_slice(&bytes[..copy_len]);
|
||||
}
|
||||
(buf, max_len)
|
||||
}
|
||||
|
||||
fn apply_compression(ds: &mut clawhdf5_format::type_builders::DatasetBuilder, opts: &WriteOptions) {
|
||||
if opts.compression {
|
||||
ds.with_deflate(opts.compression_level);
|
||||
ds.with_shuffle();
|
||||
}
|
||||
}
|
||||
|
||||
fn write_chunks_group(builder: &mut FileBuilder, data: &SqliteData, opts: &WriteOptions) {
|
||||
let mut group = builder.create_group("chunks");
|
||||
let n = data.chunks.len() as u64;
|
||||
|
||||
if n == 0 {
|
||||
group.set_attr("count", AttrValue::I64(0));
|
||||
builder.add_group(group.finish());
|
||||
return;
|
||||
}
|
||||
|
||||
group.set_attr("count", AttrValue::I64(n as i64));
|
||||
|
||||
// ids
|
||||
let ids: Vec<i64> = data.chunks.iter().map(|c| c.id).collect();
|
||||
group.create_dataset("id").with_i64_data(&ids);
|
||||
|
||||
// text
|
||||
let texts: Vec<String> = data.chunks.iter().map(|c| c.chunk.clone()).collect();
|
||||
let (text_raw, text_len) = pack_strings(&texts);
|
||||
group
|
||||
.create_dataset("text")
|
||||
.with_compound_data(string_dtype(text_len), text_raw, n);
|
||||
|
||||
// embeddings - flatten to [N, dim]
|
||||
let dim = data.embedding_dim;
|
||||
if opts.float16 {
|
||||
let f16_data: Vec<u16> = data
|
||||
.chunks
|
||||
.iter()
|
||||
.flat_map(|c| {
|
||||
c.embedding
|
||||
.iter()
|
||||
.map(|&v| half::f16::from_f32(v).to_bits())
|
||||
})
|
||||
.collect();
|
||||
let raw: Vec<u8> = f16_data.iter().flat_map(|v| v.to_le_bytes()).collect();
|
||||
let f16_dtype = Datatype::FloatingPoint {
|
||||
size: 2,
|
||||
byte_order: clawhdf5_format::datatype::DatatypeByteOrder::LittleEndian,
|
||||
bit_offset: 0,
|
||||
bit_precision: 16,
|
||||
exponent_location: 10,
|
||||
exponent_size: 5,
|
||||
mantissa_location: 0,
|
||||
mantissa_size: 10,
|
||||
exponent_bias: 15,
|
||||
};
|
||||
let ds = group
|
||||
.create_dataset("embeddings")
|
||||
.with_compound_data(f16_dtype, raw, n)
|
||||
.with_shape(&[n, dim as u64]);
|
||||
apply_compression(ds, opts);
|
||||
} else {
|
||||
let flat: Vec<f32> = data
|
||||
.chunks
|
||||
.iter()
|
||||
.flat_map(|c| c.embedding.iter().copied())
|
||||
.collect();
|
||||
let ds = group
|
||||
.create_dataset("embeddings")
|
||||
.with_f32_data(&flat)
|
||||
.with_shape(&[n, dim as u64]);
|
||||
apply_compression(ds, opts);
|
||||
}
|
||||
|
||||
// source_channel
|
||||
let channels: Vec<String> = data
|
||||
.chunks
|
||||
.iter()
|
||||
.map(|c| c.source_channel.clone())
|
||||
.collect();
|
||||
let (ch_raw, ch_len) = pack_strings(&channels);
|
||||
group
|
||||
.create_dataset("source_channel")
|
||||
.with_compound_data(string_dtype(ch_len), ch_raw, n);
|
||||
|
||||
// timestamp
|
||||
let timestamps: Vec<f64> = data.chunks.iter().map(|c| c.timestamp).collect();
|
||||
group.create_dataset("timestamp").with_f64_data(×tamps);
|
||||
|
||||
// session_id
|
||||
let sess_ids: Vec<String> = data.chunks.iter().map(|c| c.session_id.clone()).collect();
|
||||
let (sid_raw, sid_len) = pack_strings(&sess_ids);
|
||||
group
|
||||
.create_dataset("session_id")
|
||||
.with_compound_data(string_dtype(sid_len), sid_raw, n);
|
||||
|
||||
// tags
|
||||
let tags: Vec<String> = data.chunks.iter().map(|c| c.tags.clone()).collect();
|
||||
let (tag_raw, tag_len) = pack_strings(&tags);
|
||||
group
|
||||
.create_dataset("tags")
|
||||
.with_compound_data(string_dtype(tag_len), tag_raw, n);
|
||||
|
||||
// deleted
|
||||
let deleted: Vec<i32> = data.chunks.iter().map(|c| c.deleted).collect();
|
||||
group.create_dataset("deleted").with_i32_data(&deleted);
|
||||
|
||||
builder.add_group(group.finish());
|
||||
}
|
||||
|
||||
fn write_sessions_group(builder: &mut FileBuilder, data: &SqliteData) {
|
||||
let mut group = builder.create_group("sessions");
|
||||
let n = data.sessions.len() as u64;
|
||||
group.set_attr("count", AttrValue::I64(n as i64));
|
||||
|
||||
if n == 0 {
|
||||
builder.add_group(group.finish());
|
||||
return;
|
||||
}
|
||||
|
||||
let ids: Vec<String> = data.sessions.iter().map(|s| s.id.clone()).collect();
|
||||
let (id_raw, id_len) = pack_strings(&ids);
|
||||
group
|
||||
.create_dataset("id")
|
||||
.with_compound_data(string_dtype(id_len), id_raw, n);
|
||||
|
||||
let start_idxs: Vec<i64> = data.sessions.iter().map(|s| s.start_idx).collect();
|
||||
group.create_dataset("start_idx").with_i64_data(&start_idxs);
|
||||
|
||||
let end_idxs: Vec<i64> = data.sessions.iter().map(|s| s.end_idx).collect();
|
||||
group.create_dataset("end_idx").with_i64_data(&end_idxs);
|
||||
|
||||
let channels: Vec<String> = data.sessions.iter().map(|s| s.channel.clone()).collect();
|
||||
let (ch_raw, ch_len) = pack_strings(&channels);
|
||||
group
|
||||
.create_dataset("channel")
|
||||
.with_compound_data(string_dtype(ch_len), ch_raw, n);
|
||||
|
||||
let timestamps: Vec<f64> = data.sessions.iter().map(|s| s.timestamp).collect();
|
||||
group.create_dataset("timestamp").with_f64_data(×tamps);
|
||||
|
||||
let summaries: Vec<String> = data.sessions.iter().map(|s| s.summary.clone()).collect();
|
||||
let (sum_raw, sum_len) = pack_strings(&summaries);
|
||||
group
|
||||
.create_dataset("summary")
|
||||
.with_compound_data(string_dtype(sum_len), sum_raw, n);
|
||||
|
||||
builder.add_group(group.finish());
|
||||
}
|
||||
|
||||
fn write_entities_group(builder: &mut FileBuilder, data: &SqliteData) {
|
||||
let mut group = builder.create_group("entities");
|
||||
let n = data.entities.len() as u64;
|
||||
group.set_attr("count", AttrValue::I64(n as i64));
|
||||
|
||||
if n == 0 {
|
||||
builder.add_group(group.finish());
|
||||
return;
|
||||
}
|
||||
|
||||
let ids: Vec<i64> = data.entities.iter().map(|e| e.id).collect();
|
||||
group.create_dataset("id").with_i64_data(&ids);
|
||||
|
||||
let names: Vec<String> = data.entities.iter().map(|e| e.name.clone()).collect();
|
||||
let (name_raw, name_len) = pack_strings(&names);
|
||||
group
|
||||
.create_dataset("name")
|
||||
.with_compound_data(string_dtype(name_len), name_raw, n);
|
||||
|
||||
let types: Vec<String> = data
|
||||
.entities
|
||||
.iter()
|
||||
.map(|e| e.entity_type.clone())
|
||||
.collect();
|
||||
let (type_raw, type_len) = pack_strings(&types);
|
||||
group
|
||||
.create_dataset("type")
|
||||
.with_compound_data(string_dtype(type_len), type_raw, n);
|
||||
|
||||
let emb_idxs: Vec<i64> = data.entities.iter().map(|e| e.embedding_idx).collect();
|
||||
group
|
||||
.create_dataset("embedding_idx")
|
||||
.with_i64_data(&emb_idxs);
|
||||
|
||||
builder.add_group(group.finish());
|
||||
}
|
||||
|
||||
fn write_relations_group(builder: &mut FileBuilder, data: &SqliteData) {
|
||||
let mut group = builder.create_group("relations");
|
||||
let n = data.relations.len() as u64;
|
||||
group.set_attr("count", AttrValue::I64(n as i64));
|
||||
|
||||
if n == 0 {
|
||||
builder.add_group(group.finish());
|
||||
return;
|
||||
}
|
||||
|
||||
let srcs: Vec<i64> = data.relations.iter().map(|r| r.src).collect();
|
||||
group.create_dataset("src").with_i64_data(&srcs);
|
||||
|
||||
let tgts: Vec<i64> = data.relations.iter().map(|r| r.tgt).collect();
|
||||
group.create_dataset("tgt").with_i64_data(&tgts);
|
||||
|
||||
let rels: Vec<String> = data.relations.iter().map(|r| r.relation.clone()).collect();
|
||||
let (rel_raw, rel_len) = pack_strings(&rels);
|
||||
group
|
||||
.create_dataset("relation")
|
||||
.with_compound_data(string_dtype(rel_len), rel_raw, n);
|
||||
|
||||
let weights: Vec<f64> = data.relations.iter().map(|r| r.weight).collect();
|
||||
group.create_dataset("weight").with_f64_data(&weights);
|
||||
|
||||
let timestamps: Vec<f64> = data.relations.iter().map(|r| r.timestamp).collect();
|
||||
group.create_dataset("timestamp").with_f64_data(×tamps);
|
||||
|
||||
builder.add_group(group.finish());
|
||||
}
|
||||
Reference in New Issue
Block a user