lance_nd/
prune.rs

1//! Change a filter on the table columns into a filter for a scan with stats.
2//!
3//! The filter uses the Lance table columns, for example
4//! `get_field(col("sst"), "values")`, `get_field(col("sst"), "attrs")` or
5//! `get_field(col("__global_attrs__"), "id")`. lance-nd does not check or
6//! change the names, so Lance reports a wrong name in the scan.
7//!
8//! A value part reads `<variable>.values`. It changes into a filter on the
9//! stats column. The changed part keeps every row that can match, and skips
10//! a row only when the stats show that no value of the row matches. A value
11//! part without a stats form keeps every row.
12//!
13//! Each other part, for example on attributes, is exact for a row, so it
14//! stays as it is. These rules apply:
15//!
16//! - `AND`: a side without a stats form goes away, so more rows are kept.
17//! - `OR`: both sides need a form, or the full `OR` keeps every row.
18//! - `NOT`: only over parts without values. Over a value part, it keeps
19//!   every row.
20
21use arrow::datatypes::{DataType, Schema};
22use lance::deps::datafusion::common::Column;
23use lance::deps::datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
24use lance::deps::datafusion::functions::core::expr_fn::get_field;
25use lance::deps::datafusion::logical_expr::expr::{Between, BinaryExpr, InList};
26use lance::deps::datafusion::logical_expr::{Expr, Operator, col};
27use nd_arrow_array::extension::GLOBAL_ATTRIBUTES_FIELD;
28
29use crate::error::Result;
30use crate::stats::stats_column;
31
32/// A filter for a scan of `table` with the rows that can match `expr`.
33/// `None` means that every row is kept. [`crate::NdScanner::filter_expr`]
34/// uses it.
35///
36/// `expr` uses the Lance table columns, for example
37/// `get_field(col("sst"), "values")`, `get_field(col("sst"), "attrs")` or
38/// `get_field(col("__global_attrs__"), "id")`. lance-nd does not check or
39/// change the names, so Lance reports a wrong name in the scan.
40///
41/// A value part reads `<variable>.values`. It changes into a filter on the
42/// stats column. The changed part keeps every row that can match, and skips
43/// a row only when the stats show that no value of the row matches. A value
44/// part without a stats form keeps every row.
45///
46/// Each other part, for example on attributes, is exact for a row, so it
47/// stays as it is. These rules apply:
48///
49/// - `AND`: a side without a stats form goes away, so more rows are kept.
50/// - `OR`: both sides need a form, or the full `OR` keeps every row.
51/// - `NOT`: only over parts without values. Over a value part, it keeps
52///   every row.
53pub fn prune_filter(expr: &Expr, table: &Schema) -> Result<Option<Expr>> {
54    let value_type = |name: &str| {
55        let stats = table.field_with_name(&stats_column(name)).ok()?;
56        match stats.data_type() {
57            DataType::Struct(bounds) => bounds.find("min").map(|(_, f)| f.data_type().clone()),
58            _ => None,
59        }
60    };
61    Ok(rewrite(expr, &value_type))
62}
63
64/// The column and the field names of a column or a `get_field` path.
65fn field_path(expr: &Expr) -> Option<(&Column, Vec<&str>)> {
66    match expr {
67        Expr::Column(column) => Some((column, Vec::new())),
68        Expr::ScalarFunction(function) if function.name() == "get_field" => {
69            let (base, names) = function.args.split_first()?;
70            let (column, mut path) = field_path(base)?;
71            for name in names {
72                let Expr::Literal(value, _) = name else {
73                    return None;
74                };
75                path.push(value.try_as_str()??);
76            }
77            Some((column, path))
78        }
79        _ => None,
80    }
81}
82
83/// The variable of a `<variable>.values` path.
84fn values_of(expr: &Expr) -> Option<String> {
85    let (column, path) = field_path(expr)?;
86    let name = column.flat_name();
87    (name != GLOBAL_ATTRIBUTES_FIELD && path.first() == Some(&"values")).then_some(name)
88}
89
90/// `true` when `expr` reads the values of a variable.
91fn reads_values(expr: &Expr) -> bool {
92    let mut found = false;
93    expr.apply(|e| {
94        found = values_of(e).is_some();
95        Ok(if found {
96            TreeNodeRecursion::Stop
97        } else {
98            TreeNodeRecursion::Continue
99        })
100    })
101    .expect("the visit does not fail");
102    found
103}
104
105/// The stats form of `expr`, or `None` when it keeps every row.
106fn rewrite(expr: &Expr, value_type: &ValueType<'_>) -> Option<Expr> {
107    if !reads_values(expr) {
108        return Some(expr.clone());
109    }
110    match expr {
111        Expr::BinaryExpr(BinaryExpr { left, op, right }) => match op {
112            Operator::And => match (rewrite(left, value_type), rewrite(right, value_type)) {
113                (Some(l), Some(r)) => Some(l.and(r)),
114                (Some(side), None) | (None, Some(side)) => Some(side),
115                (None, None) => None,
116            },
117            Operator::Or => Some(rewrite(left, value_type)?.or(rewrite(right, value_type)?)),
118            _ => compare(left, *op, right, value_type),
119        },
120        Expr::Between(Between {
121            expr,
122            negated: false,
123            low,
124            high,
125        }) => {
126            let stats = Stats::of(expr, value_type)?;
127            let (low, high) = (stats.literal(low)?, stats.literal(high)?);
128            Some(stats.guard(stats.max().gt_eq(low).and(stats.min().lt_eq(high))))
129        }
130        Expr::InList(InList {
131            expr,
132            list,
133            negated: false,
134        }) => {
135            let stats = Stats::of(expr, value_type)?;
136            let any = list
137                .iter()
138                .map(|value| stats.literal(value).map(|v| stats.contains(v)))
139                .reduce(|a, b| Some(a?.or(b?)))??;
140            Some(stats.guard(any))
141        }
142        _ => None,
143    }
144}
145
146/// The stats form of `left op right`, with a variable on one side and a
147/// literal on the other.
148fn compare(left: &Expr, op: Operator, right: &Expr, value_type: &ValueType<'_>) -> Option<Expr> {
149    let (stats, op, value) = match (Stats::of(left, value_type), Stats::of(right, value_type)) {
150        (Some(stats), None) => {
151            let value = stats.literal(right)?;
152            (stats, op, value)
153        }
154        (None, Some(stats)) => {
155            let value = stats.literal(left)?;
156            (stats, op.swap()?, value)
157        }
158        _ => return None,
159    };
160    let condition = match op {
161        Operator::Eq => stats.contains(value),
162        Operator::Lt => stats.min().lt(value),
163        Operator::LtEq => stats.min().lt_eq(value),
164        Operator::Gt => stats.max().gt(value),
165        Operator::GtEq => stats.max().gt_eq(value),
166        _ => return None,
167    };
168    Some(stats.guard(condition))
169}
170
171/// The value type of the stats child of a variable, or `None` without one.
172type ValueType<'a> = dyn Fn(&str) -> Option<DataType> + 'a;
173
174/// `true` when Lance can compare values of `value` with a literal of `literal`.
175fn comparable(value: &DataType, literal: &DataType) -> bool {
176    use DataType::*;
177    let numeric = |t: &DataType| t.is_integer() || t.is_floating();
178    let string = |t: &DataType| matches!(t, Utf8 | LargeUtf8 | Utf8View);
179    (numeric(value) && numeric(literal))
180        || (string(value) && string(literal))
181        || matches!(
182            (value, literal),
183            (Timestamp(..), Timestamp(..))
184                | (Date32 | Date64, Date32 | Date64)
185                | (Time32(_) | Time64(_), Time32(_) | Time64(_))
186                | (Duration(_), Duration(_))
187        )
188}
189
190/// The stats child of one variable, and its value type.
191struct Stats(Expr, DataType);
192
193impl Stats {
194    /// The stats of `expr` when it reads the values of a variable with a
195    /// stats child.
196    fn of(expr: &Expr, value_type: &ValueType<'_>) -> Option<Self> {
197        let name = values_of(expr)?;
198        let value_type = value_type(&name)?;
199        let child = col(Column::new_unqualified(stats_column(&name)));
200        Some(Self(child, value_type))
201    }
202
203    /// `expr` when it is a literal that compares with the values.
204    fn literal(&self, expr: &Expr) -> Option<Expr> {
205        let Expr::Literal(value, _) = expr else {
206            return None;
207        };
208        comparable(&self.1, &value.data_type()).then(|| expr.clone())
209    }
210
211    fn min(&self) -> Expr {
212        get_field(self.0.clone(), "min")
213    }
214
215    fn max(&self) -> Expr {
216        get_field(self.0.clone(), "max")
217    }
218
219    /// The range of the row holds `value`.
220    fn contains(&self, value: Expr) -> Expr {
221        self.min().lt_eq(value.clone()).and(self.max().gt_eq(value))
222    }
223
224    /// `condition` for a row that has the variable. A row without counted
225    /// values (null min) is kept. A row without the variable is skipped.
226    fn guard(&self, condition: Expr) -> Expr {
227        self.0
228            .clone()
229            .is_not_null()
230            .and(self.min().is_null().or(condition))
231    }
232}
lance_nd/prune.rs

lance_nd/
prune.rs
