lance_nd/
stats.rs

1//! Per-row min and max values of each variable, for pruning.
2//!
3//! A table with stats has the schema metadata key [`STATS_METADATA_KEY`]. Each
4//! variable with a stats type has a stats column `__stats__:<variable>`, a
5//! struct of `min` and `max` typed by the value type. A writer adds a variable
6//! and its stats column together, so an older fragment lacks whole columns
7//! only, which Lance reads as null. A null row means that the row does not
8//! have the variable. Null and NaN values do not count. A variable without
9//! counted values has null `min` and `max`.
10
11use std::cmp::Ordering;
12use std::sync::Arc;
13
14use arrow::array::{
15    Array, ArrayRef, AsArray, PrimitiveArray, RecordBatch, StringArray, StructArray, UInt64Array,
16    downcast_primitive_array, make_comparator, new_empty_array, new_null_array,
17};
18use arrow::compute::kernels::aggregate::{
19    max, max_string, max_string_view, min, min_string, min_string_view,
20};
21use arrow::compute::{SortOptions, cast, concat, take};
22use arrow::datatypes::{
23    ArrowPrimitiveType, DataType, Field, Fields, Float16Type, Float32Type, Float64Type, Schema,
24};
25use nd_arrow_array::extension::{VALUES_FIELD, is_nd_field};
26
27use crate::error::Result;
28
29/// The name prefix of the stats columns. Each variable name that starts with
30/// it is reserved.
31pub const STATS_PREFIX: &str = "__stats__";
32
33/// The table schema metadata key that marks a table with stats.
34pub const STATS_METADATA_KEY: &str = "lance_nd:stats";
35
36/// The value of [`STATS_METADATA_KEY`] in a table with stats.
37const STATS_ON: &str = "true";
38
39/// `true` when the schema marks a table with stats.
40pub fn has_stats(schema: &Schema) -> bool {
41    schema
42        .metadata()
43        .get(STATS_METADATA_KEY)
44        .is_some_and(|value| value == STATS_ON)
45}
46
47/// The name of the stats column of `variable`.
48pub fn stats_column(variable: &str) -> String {
49    // Lance 11 does not allow `.` in a top-level name.
50    format!("{STATS_PREFIX}:{variable}")
51}
52
53/// The variable of the stats column `column`, or `None` for another name.
54pub fn stats_variable(column: &str) -> Option<&str> {
55    column.strip_prefix(STATS_PREFIX)?.strip_prefix(':')
56}
57
58/// `true` for a reserved name: a stats column or another name that starts
59/// with [`STATS_PREFIX`].
60pub fn is_stats_name(name: &str) -> bool {
61    name.starts_with(STATS_PREFIX)
62}
63
64/// `true` when values of `data_type` get a min and a max.
65pub fn is_stats_type(data_type: &DataType) -> bool {
66    data_type.is_integer()
67        || data_type.is_floating()
68        || matches!(
69            data_type,
70            DataType::Utf8
71                | DataType::LargeUtf8
72                | DataType::Utf8View
73                | DataType::Timestamp(_, _)
74                | DataType::Date32
75                | DataType::Date64
76                | DataType::Time32(_)
77                | DataType::Time64(_)
78                | DataType::Duration(_)
79        )
80}
81
82/// `batch` with a stats column for each variable of a stats type, and the
83/// stats marker in its metadata.
84///
85/// `batch` is an encoded batch of nd columns. Without a variable of a stats
86/// type, only the marker is added.
87pub fn add_stats(batch: &RecordBatch) -> Result<RecordBatch> {
88    let mut fields: Vec<Arc<Field>> = batch.schema().fields().iter().cloned().collect();
89    let mut columns = batch.columns().to_vec();
90    for (field, column) in batch.schema().fields().iter().zip(batch.columns()) {
91        if !is_nd_field(field) {
92            continue;
93        }
94        let nd = column.as_struct();
95        let values = nd
96            .column_by_name(VALUES_FIELD)
97            .expect("an nd column has values")
98            .as_list::<i32>();
99        let value_type = values.value_type();
100        if !is_stats_type(&value_type) {
101            continue;
102        }
103        let (min, max) = min_max(nd, values.offsets(), values.values().as_ref())?;
104        let bounds = Fields::from(vec![
105            Field::new("min", value_type.clone(), true),
106            Field::new("max", value_type, true),
107        ]);
108        let stats = StructArray::try_new(bounds.clone(), vec![min, max], nd.nulls().cloned())?;
109        let name = stats_column(field.name());
110        fields.push(Arc::new(Field::new(name, DataType::Struct(bounds), true)));
111        columns.push(Arc::new(stats));
112    }
113
114    let mut metadata = batch.schema().metadata().clone();
115    metadata.insert(STATS_METADATA_KEY.to_string(), STATS_ON.to_string());
116    let schema = Schema::new_with_metadata(fields, metadata);
117    Ok(RecordBatch::try_new(Arc::new(schema), columns)?)
118}
119
120/// The min and the max of each row of an nd column. An absent row and a row
121/// without counted values give nulls.
122fn min_max(nd: &StructArray, offsets: &[i32], values: &dyn Array) -> Result<(ArrayRef, ArrayRef)> {
123    let mut mins: Vec<ArrayRef> = Vec::with_capacity(nd.len());
124    let mut maxs: Vec<ArrayRef> = Vec::with_capacity(nd.len());
125    for row in 0..nd.len() {
126        let (min, max) = if nd.is_valid(row) {
127            let start = offsets[row] as usize;
128            let end = offsets[row + 1] as usize;
129            row_bounds(values.slice(start, end - start).as_ref())?
130        } else {
131            let null = new_null_array(values.data_type(), 1);
132            (null.clone(), null)
133        };
134        mins.push(min);
135        maxs.push(max);
136    }
137    let min = join(&mins, values.data_type())?;
138    let max = join(&maxs, values.data_type())?;
139    Ok((widen_zero(&min, -0.0), widen_zero(&max, 0.0)))
140}
141
142fn join(arrays: &[ArrayRef], data_type: &DataType) -> Result<ArrayRef> {
143    if arrays.is_empty() {
144        return Ok(new_empty_array(data_type));
145    }
146    let arrays: Vec<&dyn Array> = arrays.iter().map(|a| a.as_ref()).collect();
147    Ok(concat(&arrays)?)
148}
149
150/// The min and the max of the values of one row, as arrays of one value.
151fn row_bounds(values: &dyn Array) -> Result<(ArrayRef, ArrayRef)> {
152    let data_type = values.data_type().clone();
153    // The kernels skip nulls but count NaN, so a NaN result takes the slow path.
154    let fast: Option<(ArrayRef, ArrayRef)> = downcast_primitive_array!(
155        values => {
156            let (low, high) = (min(values), max(values));
157            (!is_nan(low) && !is_nan(high))
158                .then(|| (one(values, low, &data_type), one(values, high, &data_type)))
159        }
160        DataType::Utf8 => {
161            let values = values.as_string::<i32>();
162            Some(strings(min_string(values), max_string(values)))
163        }
164        DataType::LargeUtf8 => {
165            let values = values.as_string::<i64>();
166            Some(strings(min_string(values), max_string(values)))
167        }
168        DataType::Utf8View => {
169            let values = values.as_string_view();
170            Some(strings(min_string_view(values), max_string_view(values)))
171        }
172        _ => None
173    );
174    match fast {
175        Some(bounds) => {
176            let (low, high) = bounds;
177            Ok((cast(&low, &data_type)?, cast(&high, &data_type)?))
178        }
179        None => compared_bounds(values),
180    }
181}
182
183/// `true` for a NaN. Only a NaN is not equal to itself.
184#[allow(clippy::eq_op)]
185fn is_nan<T: PartialEq>(value: Option<T>) -> bool {
186    value.is_some_and(|v| v != v)
187}
188
189/// An array of `value`, of the type of `like`.
190fn one<T: ArrowPrimitiveType>(
191    _like: &PrimitiveArray<T>,
192    value: Option<T::Native>,
193    data_type: &DataType,
194) -> ArrayRef {
195    let array: PrimitiveArray<T> = [value].into_iter().collect();
196    Arc::new(array.with_data_type(data_type.clone()))
197}
198
199fn strings(low: Option<&str>, high: Option<&str>) -> (ArrayRef, ArrayRef) {
200    (
201        Arc::new(StringArray::from(vec![low])),
202        Arc::new(StringArray::from(vec![high])),
203    )
204}
205
206/// The bounds of `values` by comparison, without null and NaN values.
207fn compared_bounds(values: &dyn Array) -> Result<(ArrayRef, ArrayRef)> {
208    let compare = make_comparator(values, values, SortOptions::default())?;
209    let mut low: Option<usize> = None;
210    let mut high: Option<usize> = None;
211    for i in (0..values.len()).filter(|&i| values.is_valid(i) && !is_nan_at(values, i)) {
212        if low.is_none_or(|l| compare(i, l) == Ordering::Less) {
213            low = Some(i);
214        }
215        if high.is_none_or(|h| compare(i, h) == Ordering::Greater) {
216            high = Some(i);
217        }
218    }
219    let low = take(
220        values,
221        &UInt64Array::from(vec![low.map(|i| i as u64)]),
222        None,
223    )?;
224    let high = take(
225        values,
226        &UInt64Array::from(vec![high.map(|i| i as u64)]),
227        None,
228    )?;
229    Ok((low, high))
230}
231
232/// `bounds` with each zero as `zero`. Total order puts -0.0 below 0.0, but a
233/// filter treats them as equal, so a min gets -0.0 and a max gets 0.0.
234fn widen_zero(bounds: &ArrayRef, zero: f64) -> ArrayRef {
235    match bounds.data_type() {
236        DataType::Float16 => {
237            let zero = <Float16Type as ArrowPrimitiveType>::Native::from_f64(zero);
238            Arc::new(
239                bounds
240                    .as_primitive::<Float16Type>()
241                    .unary::<_, Float16Type>(|v| if v.to_f64() == 0.0 { zero } else { v }),
242            )
243        }
244        DataType::Float32 => Arc::new(
245            bounds
246                .as_primitive::<Float32Type>()
247                .unary::<_, Float32Type>(|v| if v == 0.0 { zero as f32 } else { v }),
248        ),
249        DataType::Float64 => Arc::new(
250            bounds
251                .as_primitive::<Float64Type>()
252                .unary::<_, Float64Type>(|v| if v == 0.0 { zero } else { v }),
253        ),
254        _ => bounds.clone(),
255    }
256}
257
258fn is_nan_at(values: &dyn Array, i: usize) -> bool {
259    match values.data_type() {
260        DataType::Float16 => values.as_primitive::<Float16Type>().value(i).is_nan(),
261        DataType::Float32 => values.as_primitive::<Float32Type>().value(i).is_nan(),
262        DataType::Float64 => values.as_primitive::<Float64Type>().value(i).is_nan(),
263        _ => false,
264    }
265}
266
267#[cfg(test)]
268mod tests {
269    use arrow::array::{
270        Array, ArrayRef, AsArray, BooleanArray, Float64Array, Int32Array, StringArray,
271        TimestampNanosecondArray,
272    };
273    use arrow::datatypes::{Float64Type, Int32Type, TimeUnit};
274    use nd_arrow_array::{Dimensions, NdArrowArray, NdRecordBatch};
275
276    use super::*;
277    use crate::encode::encode_rows;
278
279    fn row(variables: Vec<(&str, ArrayRef)>) -> NdRecordBatch {
280        let fields: Vec<Field> = variables
281            .iter()
282            .map(|(name, array)| Field::new(*name, array.data_type().clone(), true))
283            .collect();
284        let arrays = variables
285            .into_iter()
286            .map(|(name, array)| {
287                let dims = Dimensions::try_new([(format!("{name}_dim"), array.len())]).unwrap();
288                NdArrowArray::try_new(array, dims).unwrap()
289            })
290            .collect();
291        NdRecordBatch::try_new_inferred(Arc::new(Schema::new(fields)), arrays).unwrap()
292    }
293
294    fn stats_of(rows: &[NdRecordBatch]) -> RecordBatch {
295        add_stats(&encode_rows(rows).unwrap()).unwrap()
296    }
297
298    fn child<'a>(batch: &'a RecordBatch, variable: &str) -> &'a StructArray {
299        batch
300            .column_by_name(&stats_column(variable))
301            .expect("a stats column")
302            .as_struct()
303    }
304
305    /// The names of the stats columns of `schema`.
306    fn stats_names(schema: &Schema) -> Vec<String> {
307        schema
308            .fields()
309            .iter()
310            .map(|f| f.name().clone())
311            .filter(|name| is_stats_name(name))
312            .collect()
313    }
314
315    fn min_max(child: &StructArray) -> (&ArrayRef, &ArrayRef) {
316        (
317            child.column_by_name("min").unwrap(),
318            child.column_by_name("max").unwrap(),
319        )
320    }
321
322    #[test]
323    fn min_max_per_row() {
324        let batch = stats_of(&[
325            row(vec![
326                (
327                    "lat",
328                    Arc::new(Float64Array::from(vec![-90.0, -89.0, -88.0])),
329                ),
330                (
331                    "sst",
332                    Arc::new(Float64Array::from(vec![Some(5.0), None, Some(1.0)])),
333                ),
334            ]),
335            row(vec![
336                ("lat", Arc::new(Float64Array::from(vec![11.0, 10.0]))),
337                ("sst", Arc::new(Float64Array::from(vec![None, Some(7.0)]))),
338            ]),
339        ]);
340        let (min, max) = min_max(child(&batch, "lat"));
341        assert_eq!(min.as_primitive::<Float64Type>().values(), &[-90.0, 10.0]);
342        assert_eq!(max.as_primitive::<Float64Type>().values(), &[-88.0, 11.0]);
343        let (min, max) = min_max(child(&batch, "sst"));
344        assert_eq!(min.as_primitive::<Float64Type>().values(), &[1.0, 7.0]);
345        assert_eq!(max.as_primitive::<Float64Type>().values(), &[5.0, 7.0]);
346        assert!(has_stats(batch.schema().as_ref()));
347    }
348
349    #[test]
350    fn nan_and_all_null() {
351        let batch = stats_of(&[
352            row(vec![(
353                "v",
354                Arc::new(Float64Array::from(vec![f64::NAN, 2.0, f64::NAN])),
355            )]),
356            row(vec![(
357                "v",
358                Arc::new(Float64Array::from(vec![Some(f64::NAN), None])),
359            )]),
360        ]);
361        let stats = child(&batch, "v");
362        assert!(stats.is_valid(0) && stats.is_valid(1));
363        let (min, max) = min_max(stats);
364        assert_eq!(min.as_primitive::<Float64Type>().value(0), 2.0);
365        assert_eq!(max.as_primitive::<Float64Type>().value(0), 2.0);
366        assert!(min.is_null(1) && max.is_null(1));
367    }
368
369    #[test]
370    fn negative_nan_and_large_rows() {
371        let mut big: Vec<Option<f64>> = (0..10_000).map(|i| Some(i as f64)).collect();
372        big[17] = Some(-f64::NAN);
373        big[18] = None;
374        let batch = stats_of(&[
375            row(vec![(
376                "v",
377                Arc::new(Float64Array::from(vec![-f64::NAN, 1.0, 2.0])),
378            )]),
379            row(vec![("v", Arc::new(Float64Array::from(big)))]),
380        ]);
381        let (min, max) = min_max(child(&batch, "v"));
382        assert_eq!(min.as_primitive::<Float64Type>().values(), &[1.0, 0.0]);
383        assert_eq!(max.as_primitive::<Float64Type>().values(), &[2.0, 9999.0]);
384    }
385
386    #[test]
387    fn absent_variable_is_null_child() {
388        let batch = stats_of(&[
389            row(vec![(
390                "pressure",
391                Arc::new(Int32Array::from(vec![3, 1, 2])),
392            )]),
393            row(vec![("other", Arc::new(Int32Array::from(vec![0])))]),
394        ]);
395        let stats = child(&batch, "pressure");
396        assert!(stats.is_valid(0) && stats.is_null(1));
397        let (min, max) = min_max(stats);
398        assert_eq!(min.as_primitive::<Int32Type>().value(0), 1);
399        assert_eq!(max.as_primitive::<Int32Type>().value(0), 3);
400    }
401
402    #[test]
403    fn strings_and_timestamps() {
404        let times = TimestampNanosecondArray::from(vec![30, 10, 20]).with_timezone("UTC");
405        let batch = stats_of(&[row(vec![
406            ("name", Arc::new(StringArray::from(vec!["b", "a", "c"]))),
407            ("time", Arc::new(times)),
408        ])]);
409        let (min, max) = min_max(child(&batch, "name"));
410        assert_eq!(min.as_string::<i32>().value(0), "a");
411        assert_eq!(max.as_string::<i32>().value(0), "c");
412        let (min, max) = min_max(child(&batch, "time"));
413        assert_eq!(
414            min.data_type(),
415            &DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into()))
416        );
417        assert_eq!(
418            min.as_primitive::<arrow::datatypes::TimestampNanosecondType>()
419                .value(0),
420            10
421        );
422        assert_eq!(
423            max.as_primitive::<arrow::datatypes::TimestampNanosecondType>()
424                .value(0),
425            30
426        );
427    }
428
429    #[test]
430    fn other_types_get_no_child() {
431        let batch = stats_of(&[row(vec![
432            ("flag", Arc::new(BooleanArray::from(vec![true, false]))),
433            ("v", Arc::new(Int32Array::from(vec![1]))),
434        ])]);
435        assert_eq!(stats_names(&batch.schema()), ["__stats__:v"]);
436        assert_eq!(stats_column("v"), format!("{STATS_PREFIX}:v"));
437
438        let only_flags = stats_of(&[row(vec![(
439            "flag",
440            Arc::new(BooleanArray::from(vec![true])),
441        )])]);
442        assert!(stats_names(&only_flags.schema()).is_empty());
443        assert!(has_stats(only_flags.schema().as_ref()));
444    }
445}
lance_nd/stats.rs

lance_nd/
stats.rs
