lance_csr/
key_statistics.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Distinct-count statistics for join-key columns, so DataFusion sizes joins correctly.
5//!
6//! DataFusion 54 estimates an equi-join's output as `L × R / max(NDV_l, NDV_r)`.
7//! When a column reports no distinct count (true for every column of `MemTable` and
8//! Lance's `LanceTableProvider`), it substitutes the row count. For a graph join chain
9//! that turns a multi-million-row intermediate into an estimate of a few thousand
10//! rows, under the 128k-row `CollectLeft` threshold, and DataFusion then builds that
11//! intermediate as one single-threaded hash table.
12//!
13//! [`KeyStatisticsProvider`] wraps a table and reports known distinct counts for its
14//! node-id and edge-endpoint columns; everything else is the inner table's.
15//! [`KeyStatisticsExec`] is the scan-side wrapper. It is transparent like
16//! DataFusion's own `CooperativeExec`: filters, projections, limits and sort pushdown
17//! pass through to the inner scan unchanged.
18
19use std::collections::HashMap;
20use std::fmt;
21use std::sync::Arc;
22
23use arrow_schema::SchemaRef;
24use async_trait::async_trait;
25use datafusion::catalog::{ScanArgs, ScanResult, Session};
26use datafusion::common::config::ConfigOptions;
27use datafusion::common::stats::Precision;
28use datafusion::common::{Constraints, Result, Statistics};
29use datafusion::datasource::{TableProvider, TableType};
30use datafusion::execution::{SendableRecordBatchStream, TaskContext};
31use datafusion::logical_expr::{Expr, TableProviderFilterPushDown};
32use datafusion::physical_expr::PhysicalExpr;
33use datafusion::physical_plan::execution_plan::CardinalityEffect;
34use datafusion::physical_plan::filter_pushdown::{
35    ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation,
36};
37use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
38
39/// A table that reports `distinct` (column name -> distinct non-null values) as
40/// column statistics, on top of whatever the inner table reports.
41///
42/// The counts are estimates for planning only; wrong counts give worse plans, never
43/// wrong results.
44#[derive(Debug)]
45pub struct KeyStatisticsProvider {
46    inner: Arc<dyn TableProvider>,
47    distinct: HashMap<String, usize>,
48}
49
50impl KeyStatisticsProvider {
51    pub fn new(inner: Arc<dyn TableProvider>, distinct: HashMap<String, usize>) -> Self {
52        Self { inner, distinct }
53    }
54}
55
56#[async_trait]
57impl TableProvider for KeyStatisticsProvider {
58    fn schema(&self) -> SchemaRef {
59        self.inner.schema()
60    }
61
62    fn constraints(&self) -> Option<&Constraints> {
63        self.inner.constraints()
64    }
65
66    fn table_type(&self) -> TableType {
67        self.inner.table_type()
68    }
69
70    async fn scan(
71        &self,
72        state: &dyn Session,
73        projection: Option<&Vec<usize>>,
74        filters: &[Expr],
75        limit: Option<usize>,
76    ) -> Result<Arc<dyn ExecutionPlan>> {
77        let plan = self.inner.scan(state, projection, filters, limit).await?;
78        let divisor = equality_divisor(filters, &self.distinct);
79        Ok(KeyStatisticsExec::wrap(plan, &self.distinct, divisor))
80    }
81
82    async fn scan_with_args<'a>(
83        &self,
84        state: &dyn Session,
85        args: ScanArgs<'a>,
86    ) -> Result<ScanResult> {
87        let divisor = equality_divisor(args.filters().unwrap_or_default(), &self.distinct);
88        let plan = self.inner.scan_with_args(state, args).await?.into_inner();
89        Ok(ScanResult::new(KeyStatisticsExec::wrap(
90            plan,
91            &self.distinct,
92            divisor,
93        )))
94    }
95
96    fn supports_filters_pushdown(
97        &self,
98        filters: &[&Expr],
99    ) -> Result<Vec<TableProviderFilterPushDown>> {
100        self.inner.supports_filters_pushdown(filters)
101    }
102
103    fn statistics(&self) -> Option<Statistics> {
104        self.inner.statistics()
105    }
106}
107
108/// Row-count divisor for filters pushed into the scan: each `column = literal`
109/// conjunct on a column with a known distinct count divides the rows by that count,
110/// as DataFusion's `FilterExec` estimates equality. Needed because Lance applies
111/// pushed filters exactly, so no `FilterExec` is left to estimate them and the scan
112/// reports the whole table (a lookup by id looked like 10k rows instead of 1).
113/// Other filters are ignored, which overestimates.
114fn equality_divisor(filters: &[Expr], distinct: &HashMap<String, usize>) -> usize {
115    use datafusion::logical_expr::utils::split_conjunction;
116    use datafusion::logical_expr::{BinaryExpr, Operator};
117    let mut divisor = 1usize;
118    for conjunct in filters.iter().flat_map(split_conjunction) {
119        let Expr::BinaryExpr(BinaryExpr {
120            left,
121            op: Operator::Eq,
122            right,
123        }) = conjunct
124        else {
125            continue;
126        };
127        let column = match (left.as_ref(), right.as_ref()) {
128            (Expr::Column(c), Expr::Literal(..)) | (Expr::Literal(..), Expr::Column(c)) => c,
129            _ => continue,
130        };
131        if let Some(&n) = distinct.get(&column.name) {
132            divisor = divisor.saturating_mul(n.max(1));
133        }
134    }
135    divisor
136}
137
138/// Scan wrapper that overrides `distinct_count` for the configured columns.
139#[derive(Debug)]
140pub struct KeyStatisticsExec {
141    input: Arc<dyn ExecutionPlan>,
142    /// Column name -> distinct count; applied by name to whatever the input outputs,
143    /// so it stays right when DataFusion prunes or reorders the scan's columns.
144    distinct: Arc<HashMap<String, usize>>,
145    /// From [`equality_divisor`]: the input's row count is divided by this (1 = unchanged).
146    rows_divisor: usize,
147}
148
149impl KeyStatisticsExec {
150    /// Wrap `input`, or return it unchanged when it has nothing to adjust.
151    fn wrap(
152        input: Arc<dyn ExecutionPlan>,
153        distinct: &HashMap<String, usize>,
154        rows_divisor: usize,
155    ) -> Arc<dyn ExecutionPlan> {
156        let schema = input.schema();
157        let relevant: HashMap<String, usize> = distinct
158            .iter()
159            .filter(|(name, _)| schema.field_with_name(name).is_ok())
160            .map(|(name, n)| (name.clone(), *n))
161            .collect();
162        if relevant.is_empty() && rows_divisor == 1 {
163            return input;
164        }
165        Arc::new(Self {
166            input,
167            distinct: Arc::new(relevant),
168            rows_divisor,
169        })
170    }
171}
172
173impl DisplayAs for KeyStatisticsExec {
174    fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
175        let mut columns: Vec<_> = self.distinct.iter().collect();
176        columns.sort();
177        let shown: Vec<String> = columns.iter().map(|(c, n)| format!("{c}={n}")).collect();
178        write!(f, "KeyStatisticsExec: distinct=[{}]", shown.join(", "))?;
179        if self.rows_divisor > 1 {
180            write!(f, ", rows/{}", self.rows_divisor)?;
181        }
182        Ok(())
183    }
184}
185
186impl ExecutionPlan for KeyStatisticsExec {
187    fn name(&self) -> &str {
188        "KeyStatisticsExec"
189    }
190
191    fn properties(&self) -> &Arc<PlanProperties> {
192        self.input.properties()
193    }
194
195    fn maintains_input_order(&self) -> Vec<bool> {
196        vec![true]
197    }
198
199    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
200        vec![&self.input]
201    }
202
203    fn with_new_children(
204        self: Arc<Self>,
205        mut children: Vec<Arc<dyn ExecutionPlan>>,
206    ) -> Result<Arc<dyn ExecutionPlan>> {
207        Ok(Self::wrap(
208            children.swap_remove(0),
209            &self.distinct,
210            self.rows_divisor,
211        ))
212    }
213
214    fn execute(
215        &self,
216        partition: usize,
217        context: Arc<TaskContext>,
218    ) -> Result<SendableRecordBatchStream> {
219        self.input.execute(partition, context)
220    }
221
222    fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
223        let mut stats = Statistics::clone(self.input.partition_statistics(partition)?.as_ref());
224        if self.rows_divisor > 1 {
225            if let Some(&n) = stats.num_rows.get_value() {
226                stats.num_rows = Precision::Inexact(n.div_ceil(self.rows_divisor));
227                stats.total_byte_size = stats.total_byte_size.to_inexact();
228            }
229        }
230        let schema = self.input.schema();
231        for (field, column) in schema
232            .fields()
233            .iter()
234            .zip(stats.column_statistics.iter_mut())
235        {
236            if let Some(&n) = self.distinct.get(field.name()) {
237                // Inexact: pushed-down filters may remove values. DataFusion also caps a
238                // distinct count at the (estimated) row count.
239                column.distinct_count = Precision::Inexact(n);
240            }
241        }
242        Ok(Arc::new(stats))
243    }
244
245    fn supports_limit_pushdown(&self) -> bool {
246        true
247    }
248
249    fn cardinality_effect(&self) -> CardinalityEffect {
250        CardinalityEffect::Equal
251    }
252
253    fn gather_filters_for_pushdown(
254        &self,
255        _phase: FilterPushdownPhase,
256        parent_filters: Vec<Arc<dyn PhysicalExpr>>,
257        _config: &ConfigOptions,
258    ) -> Result<FilterDescription> {
259        FilterDescription::from_children(parent_filters, &self.children())
260    }
261
262    fn handle_child_pushdown_result(
263        &self,
264        _phase: FilterPushdownPhase,
265        child_pushdown_result: ChildPushdownResult,
266        _config: &ConfigOptions,
267    ) -> Result<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> {
268        Ok(FilterPushdownPropagation::if_all(child_pushdown_result))
269    }
270}
271
272#[cfg(test)]
273mod tests {
274    use super::*;
275    use arrow_array::{Int64Array, RecordBatch};
276    use arrow_schema::{DataType, Field, Schema};
277    use datafusion::datasource::MemTable;
278    use datafusion::prelude::SessionContext;
279
280    #[tokio::test]
281    async fn scan_reports_distinct_counts_for_named_columns_only() {
282        let schema = Arc::new(Schema::new(vec![
283            Field::new("src", DataType::Int64, false),
284            Field::new("dst", DataType::Int64, false),
285        ]));
286        let batch = RecordBatch::try_new(
287            schema.clone(),
288            vec![
289                Arc::new(Int64Array::from(vec![1, 1, 2, 3])),
290                Arc::new(Int64Array::from(vec![7, 8, 7, 7])),
291            ],
292        )
293        .unwrap();
294        let inner = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap());
295        let table = KeyStatisticsProvider::new(inner, HashMap::from([("dst".to_string(), 2)]));
296        let ctx = SessionContext::new();
297        let plan = table.scan(&ctx.state(), None, &[], None).await.unwrap();
298        let stats = plan.partition_statistics(None).unwrap();
299        assert_eq!(
300            stats.column_statistics[1].distinct_count,
301            Precision::Inexact(2)
302        );
303        assert_eq!(stats.column_statistics[0].distinct_count, Precision::Absent);
304
305        // Projecting the counted column away leaves the inner plan unwrapped.
306        let plan = table
307            .scan(&ctx.state(), Some(&vec![0]), &[], None)
308            .await
309            .unwrap();
310        assert_ne!(plan.name(), "KeyStatisticsExec");
311
312        // A pushed equality filter on a counted column divides the row estimate;
313        // other filters leave it alone.
314        use datafusion::prelude::{col, lit};
315        let filters = [col("dst").eq(lit(7_i64)), col("src").gt(lit(1_i64))];
316        let plan = table
317            .scan(&ctx.state(), None, &filters, None)
318            .await
319            .unwrap();
320        let stats = plan.partition_statistics(None).unwrap();
321        assert_eq!(stats.num_rows, Precision::Inexact(2));
322    }
323}
lance_csr/key_statistics.rs

lance_csr/
key_statistics.rs
