fv-streams-engine 0.6.0

//! The row shim: a batch as fv-plan's row model and back, for the row-path fallbacks (a step chain
//! the translator could not compile for a schema) and the JSON edges. TRANSITIONAL — deleted when
//! every consumer is columnar.

use std::sync::Arc;

use arrow::array::{Array, RecordBatch};
use arrow::datatypes::Schema;
use fv_streams_types::decode::{OnBadData, SchemaMode, SourceDecoder, SourcePolicy};

/// TRANSITIONAL: a batch as the row model the operators still consume. Null cells are omitted so a
/// row carries exactly the fields its message had (absent fields never become explicit nulls),
/// which keeps today's emitted JSON identical. Deleted when the operators are columnar (1d).
pub fn batch_to_rows(batch: &RecordBatch) -> Vec<fv_plan::row::Row> {
    let schema = batch.schema();
    let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
    (0..batch.num_rows())
        .map(|i| {
            let mut row = Vec::with_capacity(names.len());
            for (c, col) in batch.columns().iter().enumerate() {
                if col.is_null(i) {
                    continue;
                }
                row.push((names[c].to_string(), fv_value_datafusion::cell_to_value(col, i)));
            }
            fv_plan::row::Row(row)
        })
        .collect()
}

/// Rows back to a batch, for the row-path fallbacks (a step chain the translator could not
/// compile for a schema): each row as its JSON object through an inferred-schema decode over every
/// row, so an absent field is a null cell and the batch prints exactly as the rows would.
pub fn rows_to_batch(rows: &[fv_plan::row::Row]) -> RecordBatch {
    if rows.is_empty() {
        return RecordBatch::new_empty(Arc::new(Schema::empty()));
    }
    let msgs: Vec<(Option<String>, i64, Vec<u8>)> = rows
        .iter()
        .enumerate()
        .map(|(i, r)| {
            let obj: serde_json::Map<String, serde_json::Value> =
                r.0.iter()
                    .map(|(k, v)| (k.clone(), fv_value::value_to_json(v)))
                    .collect();
            (None, i as i64, serde_json::to_vec(&obj).unwrap_or_default())
        })
        .collect();
    let mut dec = SourceDecoder::new(SourcePolicy {
        schema: SchemaMode::Inferred { sample: msgs.len() },
        on_bad_data: OnBadData::Drop,
    });
    match dec.decode(&msgs) {
        Ok(d) => d.batch,
        Err(_) => RecordBatch::new_empty(Arc::new(Schema::empty())),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn rows_to_batch_round_trips_through_batch_to_rows() {
        let rows = vec![
            fv_plan::row::Row(vec![
                ("rid".to_string(), fv_value::Value::Str("a".into())),
                ("n".to_string(), fv_value::Value::Num(2.0)),
            ]),
            fv_plan::row::Row(vec![
                ("rid".to_string(), fv_value::Value::Str("b".into())),
                ("flag".to_string(), fv_value::Value::Bool(true)),
            ]),
        ];
        let b = rows_to_batch(&rows);
        assert_eq!(b.num_rows(), 2);
        let back = batch_to_rows(&b);
        for (r, want) in back.iter().zip(&rows) {
            for (k, v) in &want.0 {
                assert_eq!(&r.get(k), v, "field {k}");
            }
            assert_eq!(
                r.0.len(),
                want.0.len(),
                "absent fields stay absent (null cells omitted)"
            );
        }
        assert_eq!(rows_to_batch(&[]).num_rows(), 0);
    }

    #[test]
    fn the_shim_omits_null_cells() {
        let mut d = SourceDecoder::new(SourcePolicy::default());
        let msgs = vec![
            (
                Some("k0".to_string()),
                100i64,
                br#"{"type":"bid","price":9312}"#.to_vec(),
            ),
            (Some("k1".to_string()), 101, br#"{"type":"person","name":"n"}"#.to_vec()),
        ];
        let out = d.decode(&msgs).unwrap();
        let rows = batch_to_rows(&out.batch);
        assert!(!rows[1].contains("price") && rows[1].contains("name"));
        assert_eq!(rows[0].get("price"), fv_value::Value::Num(9312.0));
    }
}