fv-streams-engine 0.6.0

//! PER-BATCH HELPERS shared by the stages: the time column as event times (masked to what an
//! operator may see), the rid-keyed emission keys, fired operator rows as keyed output rows.

use super::*;

/// The batch with precomputed emission keys riding as the `__fv_key` meta column.
pub(super) fn with_key_column(batch: &arrow::array::RecordBatch, keys: &[String]) -> arrow::array::RecordBatch {
    use arrow::datatypes::{DataType, Field, Schema};
    let mut fields: Vec<Field> = batch.schema().fields().iter().map(|f| f.as_ref().clone()).collect();
    fields.push(Field::new("__fv_key", DataType::Utf8, true));
    let mut cols = batch.columns().to_vec();
    cols.push(Arc::new(arrow::array::StringArray::from(
        keys.iter().map(|k| Some(k.as_str())).collect::<Vec<_>>(),
    )));
    arrow::array::RecordBatch::try_new(Arc::new(Schema::new(fields)), cols).expect("key column matches the batch")
}

/// The rid-keyed emission keys of a batch: the row's own `rid` column when it is a string, else
/// the consumed key, else the batch's fallback key.
pub(super) fn rid_keys(
    batch: &arrow::array::RecordBatch,
    consumed: &[Option<String>],
    fallback: Option<&str>,
) -> Vec<Option<String>> {
    use arrow::array::Array;
    let rid = batch
        .column_by_name("rid")
        .and_then(|c| c.as_any().downcast_ref::<arrow::array::StringArray>());
    (0..batch.num_rows())
        .map(|i| {
            rid.filter(|c| !c.is_null(i))
                .map(|c| c.value(i).to_string())
                .or_else(|| consumed.get(i).cloned().flatten())
                .or_else(|| fallback.map(String::from))
        })
        .collect()
}

/// A partition batch masked to the rows the operator may see (offset at or above the restore
/// floor, a usable event time), with each row's event time: the `time_column` as epoch-ms, or
/// `stamp` for every row when there is no column (processing time / no time).
pub(super) fn batch_with_times(
    batch: &arrow::array::RecordBatch,
    offsets: &[i64],
    floor: Option<i64>,
    time_column: Option<&str>,
    stamp: i64,
) -> (arrow::array::RecordBatch, Vec<i64>) {
    use arrow::array::Array;
    let n = batch.num_rows();
    let times: Vec<Option<i64>> = match time_column.and_then(|c| batch.column_by_name(c)) {
        Some(col) => match arrow::compute::cast(col, &arrow::datatypes::DataType::Int64) {
            Ok(ints) => {
                let ints = ints.as_any().downcast_ref::<arrow::array::Int64Array>().expect("int64");
                (0..n)
                    .map(|i| if ints.is_null(i) { None } else { Some(ints.value(i)) })
                    .collect()
            }
            Err(_) => vec![None; n],
        },
        None if time_column.is_some() => vec![None; n],
        None => vec![Some(stamp); n],
    };
    let keep: Vec<bool> = (0..n)
        .map(|i| times[i].is_some() && floor.is_none_or(|f| offsets[i] >= f))
        .collect();
    if keep.iter().all(|k| *k) {
        return (batch.clone(), times.into_iter().map(|t| t.unwrap_or(stamp)).collect());
    }
    let mask = arrow::array::BooleanArray::from(keep.clone());
    let filtered = arrow::compute::filter_record_batch(batch, &mask).expect("filter");
    let kept_times = times
        .into_iter()
        .zip(keep)
        .filter(|(_, k)| *k)
        .map(|(t, _)| t.unwrap_or(stamp))
        .collect();
    (filtered, kept_times)
}

/// Turn fired `WindowRow`s into emittable rows: the group-by key columns + `windowStart`/`windowEnd` +
/// each aggregation's aliased value.
/// Convert an operator's fired rows to output rows WITH their deterministic emission keys, per the
/// stage's [`EmitShape`]: window rows get bounds columns + `w|start|end|group` keys; rank rows
/// get a `rank` column + `t|group|rank` keys (the rank rides last in the WindowRow key vector); ring
/// rows get `l|group` keys. One function so every emission site (advance, both flush paths) derives
/// layout + key identically.
pub(super) fn fired_to_keyed(
    fired: Vec<fv_streams_ops::WindowRow>,
    shape: EmitShape,
    group_by: &[String],
) -> Vec<(String, fv_plan::row::Row)> {
    match shape {
        EmitShape::Window => window_rows_to_rows(fired, group_by)
            .into_iter()
            .map(|r| (window_emit_key(&r, group_by), r))
            .collect(),
        EmitShape::Rank | EmitShape::Ring => {
            let with_rank = matches!(shape, EmitShape::Rank);
            fired
                .into_iter()
                .map(|w| {
                    let mut cols: Vec<(String, fv_value::Value)> =
                        Vec::with_capacity(group_by.len() + 1 + w.aggs.len());
                    let mut key_vals = w.key.into_iter();
                    for name in group_by {
                        cols.push((name.clone(), key_vals.next().unwrap_or(fv_value::Value::Null)));
                    }
                    if with_rank {
                        cols.push(("rank".into(), key_vals.next().unwrap_or(fv_value::Value::Null)));
                    }
                    cols.extend(w.aggs);
                    let row = fv_plan::row::Row(cols);
                    (rank_emit_key(&row, group_by, with_rank), row)
                })
                .collect()
        }
    }
}

pub(super) fn window_rows_to_rows(
    fired: Vec<fv_streams_ops::WindowRow>,
    group_by: &[String],
) -> Vec<fv_plan::row::Row> {
    fired
        .into_iter()
        .map(|w| {
            let mut cols: Vec<(String, fv_value::Value)> = Vec::with_capacity(group_by.len() + 2 + w.aggs.len());
            for (name, val) in group_by.iter().zip(w.key) {
                cols.push((name.clone(), val));
            }
            cols.push(("windowStart".into(), fv_value::Value::Num(w.window_start as f64)));
            cols.push(("windowEnd".into(), fv_value::Value::Num(w.window_end as f64)));
            cols.extend(w.aggs);
            fv_plan::row::Row(cols)
        })
        .collect()
}

/// Watermark ids a sharded join instance feeds to its `StreamJoin`: the two INPUT SIDES of
/// one co-located partition pair — not Kafka partition numbers.
pub(super) const JOIN_SIDE_LEFT: i32 = 0;

pub(super) const JOIN_SIDE_RIGHT: i32 = 1;

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

    #[test]
    fn batch_with_times_masks_floors_and_missing_times() {
        let schema = Arc::new(arrow::datatypes::Schema::new(vec![arrow::datatypes::Field::new(
            "ts",
            arrow::datatypes::DataType::Float64,
            true,
        )]));
        let ts: arrow::array::Float64Array = vec![Some(10.0), None, Some(30.0), Some(40.0)].into();
        let b = arrow::array::RecordBatch::try_new(schema, vec![Arc::new(ts)]).unwrap();
        let (kept, times) = batch_with_times(&b, &[1, 2, 3, 4], Some(3), Some("ts"), 0);
        assert_eq!(kept.num_rows(), 2, "row 1 has no time, rows 0–1 are below the floor");
        assert_eq!(times, vec![30, 40]);
        let (all, t) = batch_with_times(&b, &[1, 2, 3, 4], None, None, 99);
        assert_eq!((all.num_rows(), t), (4, vec![99; 4]));
    }

    #[test]
    fn window_rows_to_rows_lays_out_key_bounds_and_aggs() {
        let spec = WindowSpec {
            time_column: "ts".into(),
            ingest_time: false,
            window_ms: 1000,
            slide_ms: None,
            allowed_lateness_ms: 0,
            idle_timeout_ms: 0,
            key_by: false,
            group_by: vec!["k".into()],
            aggs: vec![fv_streams_ops::Agg::count("n")],
            trigger: fv_streams_ops::Trigger::OnWatermark,
        };
        let fired = vec![fv_streams_ops::WindowRow {
            window_start: 2000,
            window_end: 3000,
            key: vec![fv_value::Value::Str("a".into())],
            aggs: vec![("n".into(), fv_value::Value::Num(5.0))],
        }];
        let rows = window_rows_to_rows(fired, &spec.group_by);
        assert_eq!(rows.len(), 1);
        let cols = &rows[0].0;
        assert_eq!(cols[0], ("k".into(), fv_value::Value::Str("a".into())));
        assert_eq!(cols[1], ("windowStart".into(), fv_value::Value::Num(2000.0)));
        assert_eq!(cols[2], ("windowEnd".into(), fv_value::Value::Num(3000.0)));
        assert_eq!(cols[3], ("n".into(), fv_value::Value::Num(5.0)));
    }
}