Function encode_as_arrow_ipc

nodedb/src/data/executor/response_codec/arrow.rs:19–85

Encode document rows as Arrow IPC bytes for columnar transport. Returns `None` if rows are empty or schema inference fails.

(
    rows: &[(String, serde_json::Value)],
    projection: &[String],
)

///
/// Returns `None` if rows are empty or schema inference fails.
pub fn encode_as_arrow_ipc(
    rows: &[(String, serde_json::Value)],
    projection: &[String],
) -> Option<Vec<u8>> {
    if rows.is_empty() {
        return None;
    }

    let first_obj = rows[0].1.as_object()?;
    let field_names: Vec<&str> = if projection.is_empty() {
        first_obj.keys().map(|k| k.as_str()).collect()
    } else {
        projection.iter().map(|s| s.as_str()).collect()
    };

    if field_names.is_empty() {
        return None;
    }

    let mut fields = vec![Field::new("id", DataType::Utf8, false)];
    for &name in &field_names {
        let dt = first_obj
            .get(name)
            .map(infer_type)
            .unwrap_or(DataType::Utf8);
        fields.push(Field::new(name, dt, true));
    }
    let schema = Arc::new(Schema::new(fields));

    let mut ids: Vec<String> = Vec::with_capacity(rows.len());
    let mut builders: Vec<ColBuilder> = field_names
        .iter()
        .map(|&name| {
            let dt = first_obj
                .get(name)
                .map(infer_type)
                .unwrap_or(DataType::Utf8);
            ColBuilder::new(dt, rows.len())
        })
        .collect();

    for (doc_id, data) in rows {
        ids.push(doc_id.clone());
        let obj = data.as_object();
        for (i, &name) in field_names.iter().enumerate() {
            match obj.and_then(|o| o.get(name)) {
                Some(v) => builders[i].push(v),
                None => builders[i].push_null(),
            }
        }
    }

    let mut arrays: Vec<ArrayRef> = vec![Arc::new(StringArray::from(ids))];
    for b in builders {
        arrays.push(b.finish());
    }

    let batch = RecordBatch::try_new(schema.clone(), arrays).ok()?;