nodedb_array/query/
rechunk.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Re-bucket cells into a different tile-extent layout.
4//!
5//! The same schema may be re-tiled (different `tile_extents`) for a
6//! query that needs a different access pattern — e.g. a slice along
7//! one axis benefits from longer extents on that axis. Rechunk takes
8//! a single source tile and emits the (possibly multiple) target
9//! tiles that contain its cells.
10//!
11//! Both source and target schemas must share name, dim arity, attrs
12//! and dim domains; only `tile_extents` may differ. The caller is
13//! responsible for that constraint — this is a pure re-bucket.
14
15use std::collections::BTreeMap;
16
17use crate::error::ArrayResult;
18use crate::schema::ArraySchema;
19use crate::tile::layout::tile_id_for_cell;
20use crate::tile::sparse_tile::{RowKind, SparseRow, SparseTile, SparseTileBuilder};
21use crate::types::TileId;
22use crate::types::cell_value::value::CellValue;
23use crate::types::coord::value::CoordValue;
24
25/// Re-bucket every cell in `tile` according to `target_schema.tile_extents`.
26/// Returns one entry per resulting target tile, ordered by [`TileId`].
27pub fn rechunk_sparse(
28    target_schema: &ArraySchema,
29    tile: &SparseTile,
30) -> ArrayResult<Vec<(TileId, SparseTile)>> {
31    let n = tile.row_count();
32    let mut live_idx = 0usize;
33    let mut buckets: BTreeMap<TileId, SparseTileBuilder<'_>> = BTreeMap::new();
34    for row in 0..n {
35        // Sentinel rows are not re-bucketed; rechunk is a purely spatial
36        // operation on live cell data.
37        if tile.row_kind(row)? != RowKind::Live {
38            continue;
39        }
40        let attr_row = live_idx;
41        live_idx += 1;
42        let coord: Vec<CoordValue> = tile
43            .dim_dicts
44            .iter()
45            .map(|d| d.values[d.indices[row] as usize].clone())
46            .collect();
47        let attrs: Vec<CellValue> = tile
48            .attr_cols
49            .iter()
50            .map(|col| col[attr_row].clone())
51            .collect();
52        // A stored cell keeps its identity; a derived row keeps none.
53        let surrogate = tile.row_surrogate(row)?;
54        let valid_from_ms = tile.valid_from_ms.get(row).copied().unwrap_or(0);
55        let valid_until_ms = tile
56            .valid_until_ms
57            .get(row)
58            .copied()
59            .unwrap_or(nodedb_types::OPEN_UPPER);
60        let tid = tile_id_for_cell(target_schema, &coord, 0)?;
61        let entry = buckets
62            .entry(tid)
63            .or_insert_with(|| SparseTileBuilder::new(target_schema));
64        entry.push_row(SparseRow {
65            coord: &coord,
66            attrs: &attrs,
67            surrogate,
68            valid_from_ms,
69            valid_until_ms,
70            kind: crate::tile::sparse_tile::RowKind::Live,
71        })?;
72    }
73    Ok(buckets.into_iter().map(|(k, v)| (k, v.build())).collect())
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79    use crate::schema::ArraySchemaBuilder;
80    use crate::schema::attr_spec::{AttrSpec, AttrType};
81    use crate::schema::dim_spec::{DimSpec, DimType};
82    use crate::types::domain::{Domain, DomainBound};
83
84    fn schema(extents: Vec<u64>) -> ArraySchema {
85        ArraySchemaBuilder::new("g")
86            .dim(DimSpec::new(
87                "x",
88                DimType::Int64,
89                Domain::new(DomainBound::Int64(0), DomainBound::Int64(99)),
90            ))
91            .dim(DimSpec::new(
92                "y",
93                DimType::Int64,
94                Domain::new(DomainBound::Int64(0), DomainBound::Int64(99)),
95            ))
96            .attr(AttrSpec::new("v", AttrType::Int64, true))
97            .tile_extents(extents)
98            .build()
99            .unwrap()
100    }
101
102    #[test]
103    fn rechunk_splits_into_smaller_tiles() {
104        let src = schema(vec![100, 100]);
105        let dst = schema(vec![10, 10]);
106        let mut b = SparseTileBuilder::new(&src);
107        b.push(
108            &[CoordValue::Int64(5), CoordValue::Int64(5)],
109            &[CellValue::Int64(1)],
110        )
111        .unwrap();
112        b.push(
113            &[CoordValue::Int64(50), CoordValue::Int64(50)],
114            &[CellValue::Int64(2)],
115        )
116        .unwrap();
117        let big = b.build();
118        let out = rechunk_sparse(&dst, &big).unwrap();
119        assert_eq!(out.len(), 2);
120        for (_, t) in &out {
121            assert_eq!(t.nnz(), 1);
122        }
123    }
124
125    #[test]
126    fn rechunk_preserves_total_cells() {
127        let src = schema(vec![16, 16]);
128        let dst = schema(vec![4, 4]);
129        let mut b = SparseTileBuilder::new(&src);
130        for i in 0..8i64 {
131            b.push(
132                &[CoordValue::Int64(i), CoordValue::Int64(i)],
133                &[CellValue::Int64(i)],
134            )
135            .unwrap();
136        }
137        let big = b.build();
138        let out = rechunk_sparse(&dst, &big).unwrap();
139        let total: u32 = out.iter().map(|(_, t)| t.nnz()).sum();
140        assert_eq!(total, 8);
141    }
142}
nodedb_array/query/rechunk.rs

nodedb_array/query/
rechunk.rs
