nodedb_lite/engine/array/ops/
elementwise.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! `ArrayOp::Elementwise` handler for NodeDB-Lite.
4//!
5//! Loads tiles from two arrays (left, right), pairs tiles that share the
6//! same Hilbert prefix, and calls `nodedb_array::query::elementwise::elementwise`
7//! per tile-pair. Unpaired tiles are elementwise-combined with an empty tile so
8//! the outer-join null semantics propagate correctly. Results are collected as
9//! one row per output cell.
10
11use std::collections::HashMap;
12use std::sync::Arc;
13
14use nodedb_array::query::elementwise::{BinaryOp, elementwise};
15use nodedb_array::tile::sparse_tile::SparseTile;
16use nodedb_array::types::TileId;
17use nodedb_array::{SegmentReader, TilePayload};
18use nodedb_physical::physical_plan::ArrayBinaryOp;
19use nodedb_types::result::QueryResult;
20use nodedb_types::value::Value;
21
22use crate::engine::array::engine::ArrayEngineState;
23use crate::engine::array::ops::util::cell::cell_value_to_value;
24use crate::error::LiteError;
25use crate::runtime::now_millis_i64;
26use crate::storage::engine::StorageEngine;
27
28fn map_binary_op(op: ArrayBinaryOp) -> BinaryOp {
29    match op {
30        ArrayBinaryOp::Add => BinaryOp::Add,
31        ArrayBinaryOp::Sub => BinaryOp::Sub,
32        ArrayBinaryOp::Mul => BinaryOp::Mul,
33        ArrayBinaryOp::Div => BinaryOp::Div,
34    }
35}
36
37/// Collect all sparse tiles from an array's segments + memtable into a map
38/// keyed by Hilbert prefix. The system-time cutoff is `system_as_of`.
39async fn collect_tiles_for_array<S: StorageEngine>(
40    array_state: &Arc<tokio::sync::Mutex<ArrayEngineState>>,
41    storage: &Arc<S>,
42    name: &str,
43    system_as_of: i64,
44) -> Result<HashMap<u64, Vec<SparseTile>>, LiteError> {
45    let (seg_ids, schema) = {
46        let state = array_state.lock().await;
47        let arr = state
48            .arrays
49            .get(name)
50            .ok_or_else(|| LiteError::BadRequest {
51                detail: format!("array '{name}' not found"),
52            })?;
53        let seg_ids: Vec<u64> = arr.manifest.segments.iter().map(|s| s.id).collect();
54        (seg_ids, arr.schema.clone())
55    };
56
57    let mut prefix_tiles: HashMap<u64, Vec<SparseTile>> = HashMap::new();
58
59    for seg_id in &seg_ids {
60        let bytes = crate::engine::array::segments::load_segment(storage, name, *seg_id).await?;
61        let reader = SegmentReader::open(&bytes).map_err(|e| LiteError::Storage {
62            detail: format!("open segment {seg_id}: {e}"),
63        })?;
64        for idx in 0..reader.tile_count() {
65            let entry_tile_id: TileId = reader.tiles()[idx].tile_id;
66            if entry_tile_id.system_from_ms > system_as_of {
67                continue;
68            }
69            let payload = reader.read_tile(idx).map_err(|e| LiteError::Storage {
70                detail: format!("read_tile seg {seg_id} idx {idx}: {e}"),
71            })?;
72            if let TilePayload::Sparse(tile) = payload {
73                prefix_tiles
74                    .entry(entry_tile_id.hilbert_prefix)
75                    .or_default()
76                    .push(tile);
77            }
78        }
79    }
80
81    // Memtable.
82    {
83        let mut state = array_state.lock().await;
84        let arr = state
85            .arrays
86            .get_mut(name)
87            .ok_or_else(|| LiteError::BadRequest {
88                detail: format!("array '{name}' not found"),
89            })?;
90        let mem_tiles = arr
91            .memtable
92            .drain_all_tiles_read_only(system_as_of, &schema)
93            .map_err(|e| LiteError::Storage {
94                detail: format!("memtable drain '{name}': {e}"),
95            })?;
96        for (tile_id, tile) in mem_tiles {
97            prefix_tiles
98                .entry(tile_id.hilbert_prefix)
99                .or_default()
100                .push(tile);
101        }
102    }
103
104    Ok(prefix_tiles)
105}
106
107/// Execute `ArrayOp::Elementwise` for the Lite engine.
108pub async fn elementwise_op<S: StorageEngine>(
109    array_state: &Arc<tokio::sync::Mutex<ArrayEngineState>>,
110    storage: &Arc<S>,
111    left_name: &str,
112    right_name: &str,
113    op: ArrayBinaryOp,
114) -> Result<QueryResult, LiteError> {
115    let system_as_of = now_millis_i64();
116    let binary_op = map_binary_op(op);
117
118    let schema = {
119        let state = array_state.lock().await;
120        let arr = state
121            .arrays
122            .get(left_name)
123            .ok_or_else(|| LiteError::BadRequest {
124                detail: format!("array '{left_name}' not found"),
125            })?;
126        arr.schema.clone()
127    };
128
129    let left_tiles = collect_tiles_for_array(array_state, storage, left_name, system_as_of).await?;
130    let right_tiles =
131        collect_tiles_for_array(array_state, storage, right_name, system_as_of).await?;
132
133    // Union of all Hilbert prefixes from both sides.
134    let mut all_prefixes: std::collections::HashSet<u64> = std::collections::HashSet::new();
135    all_prefixes.extend(left_tiles.keys().copied());
136    all_prefixes.extend(right_tiles.keys().copied());
137
138    let columns = vec![
139        "attrs".to_string(),
140        "valid_from_ms".to_string(),
141        "valid_until_ms".to_string(),
142    ];
143    let mut rows: Vec<Vec<Value>> = Vec::new();
144
145    let empty_tile = nodedb_array::tile::sparse_tile::SparseTileBuilder::new(&schema).build();
146
147    for prefix in all_prefixes {
148        let left_set = left_tiles.get(&prefix);
149        let right_set = right_tiles.get(&prefix);
150
151        // Merge all tiles within the same prefix on each side before combining.
152        let left_merged = merge_prefix_tiles(left_set, &schema, binary_op)?;
153        let right_merged = merge_prefix_tiles(right_set, &schema, binary_op)?;
154
155        let l = left_merged.as_ref().unwrap_or(&empty_tile);
156        let r = right_merged.as_ref().unwrap_or(&empty_tile);
157
158        let out = elementwise(&schema, l, r, binary_op).map_err(|e| LiteError::Storage {
159            detail: format!("elementwise prefix {prefix}: {e}"),
160        })?;
161
162        let n = out.nnz() as usize;
163        let attr_count = out.attr_cols.len();
164        for row in 0..n {
165            let attrs: Vec<Value> = (0..attr_count)
166                .map(|ai| {
167                    out.attr_cols
168                        .get(ai)
169                        .and_then(|col| col.get(row))
170                        .map(|cv| cell_value_to_value(cv.clone()))
171                        .unwrap_or(Value::Null)
172                })
173                .collect();
174            rows.push(vec![
175                Value::Array(attrs),
176                Value::Integer(0),
177                Value::Integer(i64::MAX),
178            ]);
179        }
180    }
181
182    Ok(QueryResult {
183        columns,
184        rows,
185        rows_affected: 0,
186        command: None,
187    })
188}
189
190/// Merge a slice of tiles on the same side (left or right) into a single
191/// tile using the elementwise op (same schema, so merging is self-consistent).
192/// Returns `None` when the input is empty or absent.
193fn merge_prefix_tiles(
194    tiles: Option<&Vec<SparseTile>>,
195    schema: &nodedb_array::schema::ArraySchema,
196    op: BinaryOp,
197) -> Result<Option<SparseTile>, LiteError> {
198    let tiles = match tiles {
199        Some(t) if !t.is_empty() => t,
200        _ => return Ok(None),
201    };
202    let mut acc = tiles[0].clone();
203    for tile in &tiles[1..] {
204        acc = elementwise(schema, &acc, tile, op).map_err(|e| LiteError::Storage {
205            detail: format!("merge prefix tiles: {e}"),
206        })?;
207    }
208    Ok(Some(acc))
209}
nodedb_lite/engine/array/ops/elementwise.rs

nodedb_lite/engine/array/ops/
elementwise.rs
