linkrs_query/executor/streaming/operators/blocking/
materialize.rs

1use std::collections::HashSet;
2use std::sync::Arc;
3
4use crate::executor::streaming::slot::SlotLayout;
5use crate::executor::streaming::spill::{HashPartitionSpiller, RunReader, SpilledRun};
6use linkrs_core::columnar::{MaterializedBatch, RowKey};
7use linkrs_core::Value;
8
9#[derive(Debug)]
10pub struct DistinctState {
11    /// Dedup keys: the key-column projection of each unique row. Empty
12    /// `key_cols` means the full row, preserving the exact `Value` Hash/Eq
13    /// contract (Null, NaN, mixed kinds) of the former `HashSet<Vec<Value>>`.
14    pub seen_rows: HashSet<RowKey>,
15    /// Key column positions. Empty selects the full row. Reserved for future
16    /// keyed distinct pushdown; no producer sets it yet.
17    pub key_cols: Vec<usize>,
18    /// Output buffer built once from drained keys at emission time, so the
19    /// outlet serves batch slices instead of a drained row iterator.
20    pub batch: MaterializedBatch,
21    /// Output cursor into `batch`.
22    pub emitted_offset: usize,
23    pub col_names: Vec<String>,
24    pub input_layout: Option<Arc<SlotLayout>>,
25    pub partition_spiller: Option<HashPartitionSpiller>,
26    pub spilled_runs: Vec<Option<SpilledRun>>,
27    pub current_partition: usize,
28    pub partition_seen: HashSet<RowKey>,
29    pub has_spilled: bool,
30}
31
32/// Project the dedup key of a row.
33///
34/// Empty `key_cols` clones the full row (identity projection, the only mode
35/// with a producer today). A future keyed distinct sets column positions to
36/// dedup by key while emitting full rows.
37pub(crate) fn project_key(row: &[Value], key_cols: &[usize]) -> RowKey {
38    if key_cols.is_empty() {
39        row.to_vec()
40    } else {
41        key_cols.iter().map(|&c| row[c].clone()).collect()
42    }
43}
44
45#[cfg(test)]
46mod tests {
47    use super::*;
48    use linkrs_core::value::NullType;
49
50    fn sample_rows() -> Vec<Vec<Value>> {
51        vec![
52            vec![Value::BigInt(1), Value::string("k")],
53            vec![Value::BigInt(1), Value::string("k")],
54            vec![Value::BigInt(2), Value::string("k")],
55            vec![Value::Null(NullType::Null), Value::string("k")],
56            vec![Value::Null(NullType::Null), Value::string("k")],
57            vec![Value::Double(f64::NAN), Value::Int(1)],
58            vec![Value::Int(1), Value::BigInt(1)],
59            vec![Value::Int(1), Value::BigInt(1)],
60        ]
61    }
62
63    fn dedup_reference(rows: &[Vec<Value>]) -> Vec<Vec<Value>> {
64        let mut seen: HashSet<Vec<Value>> = HashSet::new();
65        let mut out = Vec::new();
66        for row in rows {
67            if seen.insert(row.clone()) {
68                out.push(row.clone());
69            }
70        }
71        out
72    }
73
74    fn dedup_keyed(rows: &[Vec<Value>], key_cols: &[usize]) -> Vec<Vec<Value>> {
75        let mut seen: HashSet<RowKey> = HashSet::new();
76        let mut batch = MaterializedBatch::new(0, 0);
77        for row in rows {
78            if seen.insert(project_key(row, key_cols)) {
79                batch.append_row(row.clone());
80            }
81        }
82        batch.to_rows()
83    }
84
85    #[test]
86    fn identity_projection_matches_row_reference() {
87        let rows = sample_rows();
88        let mut keyed = dedup_keyed(&rows, &[]);
89        let mut reference = dedup_reference(&rows);
90        // Both hold the same multiset; HashSet iteration order is arbitrary.
91        keyed.sort_by(|a, b| format!("{a:?}").cmp(&format!("{b:?}")));
92        reference.sort_by(|a, b| format!("{a:?}").cmp(&format!("{b:?}")));
93        // NaN never equals itself under `Value` equality, so both keep every
94        // NaN row; compare lengths and non-NaN members instead of full eq.
95        assert_eq!(keyed.len(), reference.len());
96        for row in &reference {
97            assert!(
98                keyed.iter().any(|k| k == row),
99                "reference row missing from keyed output: {row:?}"
100            );
101        }
102    }
103
104    #[test]
105    fn key_subset_dedups_by_key_columns_only() {
106        let rows = sample_rows();
107        let keyed = dedup_keyed(&rows, &[1]);
108        // Column 1 values: "k" x5, Int(1) x1, BigInt(1) x2 → 3 groups.
109        assert_eq!(keyed.len(), 3);
110        assert_eq!(keyed[0], rows[0]);
111        // Full-row mode agrees with the row reference implementation.
112        assert_eq!(dedup_keyed(&rows, &[]).len(), dedup_reference(&rows).len());
113    }
114
115    #[test]
116    fn project_key_empty_cols_is_identity() {
117        let row = vec![Value::BigInt(1), Value::Null(NullType::Null)];
118        assert_eq!(project_key(&row, &[]), row);
119        assert_eq!(project_key(&row, &[1]), vec![Value::Null(NullType::Null)]);
120    }
121}
122
123#[derive(Debug)]
124pub struct MaterializeState {
125    /// Column-oriented materialized state (was `Vec<Vec<Value>>`).
126    pub batch: MaterializedBatch,
127    /// Output cursor into `batch`; replaces the drained row iterator so the
128    /// batch stays intact for future spill replay.
129    pub emitted_offset: usize,
130    pub materialized: bool,
131    pub input_layout: Option<Arc<SlotLayout>>,
132    /// Ordered spill runs (creation order = input order) plus streaming
133    /// replay cursor. Runs replay before the in-memory tail, preserving
134    /// input order end to end.
135    pub spilled_runs: Vec<SpilledRun>,
136    pub replay_index: usize,
137    pub replay_reader: Option<RunReader>,
138    pub replay_batch: Option<MaterializedBatch>,
139    pub replay_offset: usize,
140    /// Memory bytes currently accounted for `batch` rows.
141    pub accounted_bytes: usize,
142}
143
144#[derive(Debug)]
145pub struct DataCollectState {
146    /// Column-oriented collected state (was `Vec<Vec<Value>>`).
147    pub batch: MaterializedBatch,
148    pub emitted: bool,
149    /// Output cursor into `batch` (spill path only; the fast path emits one chunk).
150    pub emitted_offset: usize,
151    pub input_layout: Option<Arc<SlotLayout>>,
152    /// Ordered spill runs plus streaming replay cursor (see `MaterializeState`).
153    pub spilled_runs: Vec<SpilledRun>,
154    pub replay_index: usize,
155    pub replay_reader: Option<RunReader>,
156    pub replay_batch: Option<MaterializedBatch>,
157    pub replay_offset: usize,
158    /// Memory bytes currently accounted for `batch` rows.
159    pub accounted_bytes: usize,
160}
161
162#[derive(Debug)]
163pub struct RollUpApplyState {
164    /// Column-oriented accumulated state (was `Vec<Vec<Value>>`).
165    pub batch: MaterializedBatch,
166    /// Output cursor into `batch` (same contract as `MaterializeState`).
167    pub emitted_offset: usize,
168    /// True once the input has been fully accumulated.
169    pub accumulated: bool,
170    /// Ordered spill runs plus streaming replay cursor (see `MaterializeState`).
171    pub spilled_runs: Vec<SpilledRun>,
172    pub replay_index: usize,
173    pub replay_reader: Option<RunReader>,
174    pub replay_batch: Option<MaterializedBatch>,
175    pub replay_offset: usize,
176    /// Memory bytes currently accounted for `batch` rows.
177    pub accounted_bytes: usize,
178}
linkrs_query/executor/streaming/operators/blocking/materialize.rs

linkrs_query/executor/streaming/operators/blocking/
materialize.rs
