zarr_datafusion/reader/
coord.rs

1//! Coordinate utilities for Zarr array flattening
2//!
3//! Provides utilities for building DictionaryArrays from coordinate values,
4//! calculating subset ranges for limit optimization, and building coordinate keys.
5
6use arrow::array::{ArrayRef, DictionaryArray, Float32Array, Float64Array, Int16Array, Int64Array};
7use arrow::datatypes::Int16Type;
8use std::sync::Arc;
9
10/// Coordinate values that can be either i64 or f32/f64
11#[derive(Debug)]
12pub enum CoordValues {
13    Int64(Vec<i64>),
14    Float32(Vec<f32>),
15    Float64(Vec<f64>),
16}
17
18impl CoordValues {
19    pub fn len(&self) -> usize {
20        match self {
21            CoordValues::Int64(v) => v.len(),
22            CoordValues::Float32(v) => v.len(),
23            CoordValues::Float64(v) => v.len(),
24        }
25    }
26
27    #[allow(dead_code)]
28    pub fn is_empty(&self) -> bool {
29        self.len() == 0
30    }
31}
32
33/// Create a DictionaryArray for a coordinate column with proper type
34pub fn create_coord_dictionary_typed(
35    values: &CoordValues,
36    coord_idx: usize,
37    coord_sizes: &[usize],
38    total_rows: usize,
39) -> ArrayRef {
40    let keys = build_coord_keys(values.len(), coord_idx, coord_sizes, total_rows);
41    let keys_array = Int16Array::from(keys);
42
43    match values {
44        CoordValues::Int64(vals) => {
45            let values_array = Int64Array::from(vals.clone());
46            Arc::new(DictionaryArray::<Int16Type>::new(
47                keys_array,
48                Arc::new(values_array),
49            ))
50        }
51        CoordValues::Float32(vals) => {
52            let values_array = Float32Array::from(vals.clone());
53            Arc::new(DictionaryArray::<Int16Type>::new(
54                keys_array,
55                Arc::new(values_array),
56            ))
57        }
58        CoordValues::Float64(vals) => {
59            let values_array = Float64Array::from(vals.clone());
60            Arc::new(DictionaryArray::<Int16Type>::new(
61                keys_array,
62                Arc::new(values_array),
63            ))
64        }
65    }
66}
67
68/// Build keys array for DictionaryArray
69///
70/// Instead of expanding [0,1,2] to [0,0,0,1,1,1,2,2,2,...] (700 i64 values = 5600 bytes),
71/// we store:
72///   - values: [0,1,2] (the unique coordinate values)
73///   - keys: [0,0,0,1,1,1,2,2,2,...] (indices into values, as i16)
74///
75/// Memory: 700 i16 keys (1400 bytes) + 3 i64 values (24 bytes) = 1424 bytes (~75% savings)
76///
77/// References:
78/// - Arrow DictionaryArray: https://docs.rs/arrow/latest/arrow/array/struct.DictionaryArray.html
79/// - DataFusion Dictionary support: https://datafusion.apache.org/user-guide/sql/data_types.html
80pub fn build_coord_keys(
81    num_values: usize,
82    coord_idx: usize,
83    coord_sizes: &[usize],
84    total_rows: usize,
85) -> Vec<i16> {
86    let mut keys: Vec<i16> = Vec::with_capacity(total_rows);
87
88    // Elements after this coordinate (inner loop size)
89    let inner_size: usize = coord_sizes[coord_idx + 1..].iter().product();
90    let inner_size = if inner_size == 0 { 1 } else { inner_size };
91
92    // Elements before this coordinate (outer loop count)
93    let outer_count: usize = coord_sizes[..coord_idx].iter().product();
94    let outer_count = if outer_count == 0 { 1 } else { outer_count };
95
96    for _ in 0..outer_count {
97        for i in 0..num_values {
98            for _ in 0..inner_size {
99                keys.push(i as i16);
100            }
101        }
102    }
103
104    keys
105}
106
107/// Calculate the subset ranges needed for a limited number of rows
108///
109/// For row-major (C) order, the last dimension varies fastest.
110/// Given shape [a, b, c, d] and limit N, we need to figure out
111/// which ranges to read to get exactly N elements.
112pub fn calculate_limited_subset(shape: &[u64], limit: usize) -> Vec<std::ops::Range<u64>> {
113    let limit = limit as u64;
114    let mut ranges = Vec::with_capacity(shape.len());
115
116    // Work backwards from the last dimension
117    for (i, &dim_size) in shape.iter().enumerate().rev() {
118        if i == shape.len() - 1 {
119            // Last dimension: take min(limit, dim_size)
120            let take = limit.min(dim_size);
121            ranges.push(0..take);
122        } else {
123            // Earlier dimensions: calculate how many complete "slices" we need
124            let inner_size: u64 = shape[i + 1..].iter().product();
125            let slices_needed = limit.div_ceil(inner_size);
126            let take = slices_needed.min(dim_size);
127            ranges.push(0..take);
128        }
129    }
130
131    ranges.reverse();
132    ranges
133}
134
135/// Calculate how many values we need from each coordinate for a given row limit
136///
137/// For coords with sizes [a, b, c, d] and limit N rows:
138/// - coord d (last): need min(N, d) values
139/// - coord c: need ceil(N / d) values, capped at c
140/// - coord b: need ceil(N / (c*d)) values, capped at b
141/// - coord a: need ceil(N / (b*c*d)) values, capped at a
142pub fn calculate_coord_limits(coord_sizes: &[usize], limit: usize) -> Vec<usize> {
143    let mut limits = Vec::with_capacity(coord_sizes.len());
144    let n = coord_sizes.len();
145
146    for i in 0..n {
147        // Product of all coords after this one
148        let inner_size: usize = coord_sizes[i + 1..].iter().product();
149        let inner_size = if inner_size == 0 { 1 } else { inner_size };
150
151        // How many values do we need from this coord?
152        let needed = limit.div_ceil(inner_size);
153        let take = needed.min(coord_sizes[i]);
154        limits.push(take);
155    }
156
157    limits
158}
zarr_datafusion/reader/coord.rs

zarr_datafusion/reader/
coord.rs
