paimon_datafusion/physical_plan/
scan.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::any::Any;
19use std::sync::Arc;
20
21use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef;
22use datafusion::common::stats::Precision;
23use datafusion::common::Statistics;
24use datafusion::error::Result as DFResult;
25use datafusion::execution::{SendableRecordBatchStream, TaskContext};
26use datafusion::physical_expr::EquivalenceProperties;
27use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
28use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
29use datafusion::physical_plan::{DisplayAs, ExecutionPlan, Partitioning, PlanProperties};
30use futures::{StreamExt, TryStreamExt};
31use paimon::spec::Predicate;
32use paimon::table::Table;
33use paimon::DataSplit;
34
35use crate::error::to_datafusion_error;
36
37/// Execution plan that scans a Paimon table with optional column projection.
38///
39/// Planning is performed eagerly in [`super::super::table::PaimonTableProvider::scan`],
40/// and the resulting splits are distributed across DataFusion execution partitions
41/// so that DataFusion can schedule them in parallel.
42#[derive(Debug)]
43pub struct PaimonTableScan {
44    table: Table,
45    /// Projected column names (if None, reads all columns).
46    projected_columns: Option<Vec<String>>,
47    /// Filter translated from DataFusion expressions and reused during execute()
48    /// so reader-side pruning reaches the actual read path.
49    pushed_predicate: Option<Predicate>,
50    /// Pre-planned partition assignments: `planned_partitions[i]` contains the
51    /// Paimon splits that DataFusion partition `i` will read.
52    /// Wrapped in `Arc` to avoid deep-cloning `DataSplit` metadata in `execute()`.
53    planned_partitions: Vec<Arc<[DataSplit]>>,
54    plan_properties: Arc<PlanProperties>,
55    /// Optional limit hint pushed to paimon-core planning.
56    limit: Option<usize>,
57    /// Whether the pushed predicate is exact (no residual filtering needed).
58    /// When true and all splits have known merged_row_count, statistics can be exact.
59    filter_exact: bool,
60}
61
62impl PaimonTableScan {
63    pub(crate) fn new(
64        schema: ArrowSchemaRef,
65        table: Table,
66        projected_columns: Option<Vec<String>>,
67        pushed_predicate: Option<Predicate>,
68        planned_partitions: Vec<Arc<[DataSplit]>>,
69        limit: Option<usize>,
70        filter_exact: bool,
71    ) -> Self {
72        let plan_properties = Arc::new(PlanProperties::new(
73            EquivalenceProperties::new(schema.clone()),
74            Partitioning::UnknownPartitioning(planned_partitions.len()),
75            EmissionType::Incremental,
76            Boundedness::Bounded,
77        ));
78        Self {
79            table,
80            projected_columns,
81            pushed_predicate,
82            planned_partitions,
83            plan_properties,
84            limit,
85            filter_exact,
86        }
87    }
88
89    pub fn table(&self) -> &Table {
90        &self.table
91    }
92
93    #[cfg(test)]
94    pub(crate) fn planned_partitions(&self) -> &[Arc<[DataSplit]>] {
95        &self.planned_partitions
96    }
97
98    #[cfg(test)]
99    pub(crate) fn pushed_predicate(&self) -> Option<&Predicate> {
100        self.pushed_predicate.as_ref()
101    }
102
103    pub fn limit(&self) -> Option<usize> {
104        self.limit
105    }
106}
107
108impl ExecutionPlan for PaimonTableScan {
109    fn name(&self) -> &str {
110        "PaimonTableScan"
111    }
112
113    fn as_any(&self) -> &dyn Any {
114        self
115    }
116
117    fn properties(&self) -> &Arc<PlanProperties> {
118        &self.plan_properties
119    }
120
121    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan + 'static>> {
122        vec![]
123    }
124
125    fn with_new_children(
126        self: Arc<Self>,
127        _children: Vec<Arc<dyn ExecutionPlan>>,
128    ) -> DFResult<Arc<dyn ExecutionPlan>> {
129        Ok(self)
130    }
131
132    fn execute(
133        &self,
134        partition: usize,
135        _context: Arc<TaskContext>,
136    ) -> DFResult<SendableRecordBatchStream> {
137        let splits = Arc::clone(self.planned_partitions.get(partition).ok_or_else(|| {
138            datafusion::error::DataFusionError::Internal(format!(
139                "PaimonTableScan: partition index {partition} out of range (total {})",
140                self.planned_partitions.len()
141            ))
142        })?);
143
144        let table = self.table.clone();
145        let schema = self.schema();
146        let projected_columns = self.projected_columns.clone();
147        let pushed_predicate = self.pushed_predicate.clone();
148
149        let fut = async move {
150            let mut read_builder = table.new_read_builder();
151
152            if let Some(ref columns) = projected_columns {
153                let col_refs: Vec<&str> = columns.iter().map(|s| s.as_str()).collect();
154                read_builder.with_projection(&col_refs);
155            }
156            if let Some(filter) = pushed_predicate {
157                read_builder.with_filter(filter);
158            }
159
160            let read = read_builder.new_read().map_err(to_datafusion_error)?;
161            let stream = read.to_arrow(&splits).map_err(to_datafusion_error)?;
162            let stream = stream.map(|r| r.map_err(to_datafusion_error));
163
164            Ok::<_, datafusion::error::DataFusionError>(RecordBatchStreamAdapter::new(
165                schema,
166                Box::pin(stream),
167            ))
168        };
169
170        Ok(Box::pin(RecordBatchStreamAdapter::new(
171            self.schema(),
172            futures::stream::once(fut).try_flatten(),
173        )))
174    }
175
176    fn partition_statistics(&self, partition: Option<usize>) -> DFResult<Statistics> {
177        let partitions: &[Arc<[DataSplit]>] = match partition {
178            Some(idx) => std::slice::from_ref(&self.planned_partitions[idx]),
179            None => &self.planned_partitions,
180        };
181
182        let mut total_rows: usize = 0;
183        let mut total_bytes: usize = 0;
184        let mut all_row_counts_known = true;
185        for splits in partitions {
186            for split in splits.iter() {
187                if let Some(row_count) = split.merged_row_count() {
188                    total_rows += row_count as usize;
189                } else {
190                    all_row_counts_known = false;
191                    total_rows += split.row_count() as usize;
192                }
193                for file in split.data_files() {
194                    total_bytes += file.file_size as usize;
195                }
196            }
197        }
198
199        // Return exact statistics when:
200        // 1. All splits have known merged_row_count (no deletion files with unknown cardinality)
201        // 2. No limit is applied (limit would make row count inexact)
202        // 3. Filter is exact (no residual filtering needed above the scan)
203        let num_rows_precision =
204            if all_row_counts_known && self.limit.is_none() && self.filter_exact {
205                Precision::Exact(total_rows)
206            } else {
207                Precision::Inexact(total_rows)
208            };
209
210        Ok(Statistics {
211            num_rows: num_rows_precision,
212            total_byte_size: Precision::Inexact(total_bytes),
213            column_statistics: Statistics::unknown_column(&self.schema()),
214        })
215    }
216}
217
218impl DisplayAs for PaimonTableScan {
219    fn fmt_as(
220        &self,
221        _t: datafusion::physical_plan::DisplayFormatType,
222        f: &mut std::fmt::Formatter,
223    ) -> std::fmt::Result {
224        write!(f, "PaimonTableScan: table={}", self.table.identifier())?;
225
226        let total_splits: usize = self.planned_partitions.iter().map(|p| p.len()).sum();
227        let total_files: usize = self
228            .planned_partitions
229            .iter()
230            .flat_map(|p| p.iter())
231            .map(|s| s.data_files().len())
232            .sum();
233        write!(
234            f,
235            ", partitions={}, splits={total_splits}, files={total_files}",
236            self.planned_partitions.len()
237        )?;
238
239        if let Some(ref columns) = self.projected_columns {
240            write!(f, ", projection=[{}]", columns.join(", "))?;
241        }
242        if let Some(ref predicate) = self.pushed_predicate {
243            write!(f, ", predicate={predicate}")?;
244        }
245        if let Some(limit) = self.limit {
246            write!(f, ", limit={limit}")?;
247        }
248        Ok(())
249    }
250}
251
252#[cfg(test)]
253mod tests {
254    use super::*;
255    mod test_utils {
256        include!(concat!(env!("CARGO_MANIFEST_DIR"), "/../../test_utils.rs"));
257    }
258
259    use datafusion::arrow::array::Int32Array;
260    use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema};
261    use datafusion::physical_plan::ExecutionPlan;
262    use datafusion::prelude::SessionContext;
263    use futures::TryStreamExt;
264    use paimon::catalog::Identifier;
265    use paimon::io::FileIOBuilder;
266    use paimon::spec::{
267        BinaryRow, DataType, Datum, IntType, PredicateBuilder, Schema as PaimonSchema, TableSchema,
268    };
269    use paimon::table::Table;
270    use std::fs;
271    use tempfile::tempdir;
272    use test_utils::{local_file_path, test_data_file, write_int_parquet_file};
273
274    fn test_schema() -> ArrowSchemaRef {
275        Arc::new(ArrowSchema::new(vec![Field::new(
276            "id",
277            ArrowDataType::Int32,
278            false,
279        )]))
280    }
281
282    #[test]
283    fn test_partition_count_empty_plan() {
284        let schema = test_schema();
285        let scan = PaimonTableScan::new(
286            schema,
287            dummy_table(),
288            None,
289            None,
290            vec![Arc::from(Vec::new())],
291            None,
292            false,
293        );
294        assert_eq!(scan.properties().output_partitioning().partition_count(), 1);
295    }
296
297    #[test]
298    fn test_partition_count_multiple_partitions() {
299        let schema = test_schema();
300        let planned_partitions = vec![
301            Arc::from(Vec::new()),
302            Arc::from(Vec::new()),
303            Arc::from(Vec::new()),
304        ];
305        let scan = PaimonTableScan::new(
306            schema,
307            dummy_table(),
308            None,
309            None,
310            planned_partitions,
311            None,
312            false,
313        );
314        assert_eq!(scan.properties().output_partitioning().partition_count(), 3);
315    }
316
317    /// Constructs a minimal Table for testing (no real files needed since we
318    /// only test PlanProperties, not actual reads).
319    fn dummy_table() -> Table {
320        let file_io = FileIOBuilder::new("file").build().unwrap();
321        let schema = PaimonSchema::builder().build().unwrap();
322        let table_schema = TableSchema::new(0, &schema);
323        Table::new(
324            file_io,
325            Identifier::new("test_db", "test_table"),
326            "/tmp/test-table".to_string(),
327            table_schema,
328            None,
329        )
330    }
331
332    #[tokio::test]
333    async fn test_execute_applies_pushed_filter_during_read() {
334        let tempdir = tempdir().unwrap();
335        let table_path = local_file_path(tempdir.path());
336        let bucket_dir = tempdir.path().join("bucket-0");
337        fs::create_dir_all(&bucket_dir).unwrap();
338
339        write_int_parquet_file(
340            &bucket_dir.join("data.parquet"),
341            vec![("id", vec![1, 2, 3, 4]), ("value", vec![5, 20, 30, 40])],
342            Some(2),
343        );
344        let file_size = fs::metadata(bucket_dir.join("data.parquet")).unwrap().len() as i64;
345
346        let file_io = FileIOBuilder::new("file").build().unwrap();
347        let table_schema = TableSchema::new(
348            0,
349            &paimon::spec::Schema::builder()
350                .column("id", DataType::Int(IntType::new()))
351                .column("value", DataType::Int(IntType::new()))
352                .build()
353                .unwrap(),
354        );
355        let table = Table::new(
356            file_io,
357            Identifier::new("default", "t"),
358            table_path,
359            table_schema,
360            None,
361        );
362
363        let split = paimon::DataSplitBuilder::new()
364            .with_snapshot(1)
365            .with_partition(BinaryRow::new(0))
366            .with_bucket(0)
367            .with_bucket_path(local_file_path(&bucket_dir))
368            .with_total_buckets(1)
369            .with_data_files(vec![test_data_file("data.parquet", 4, file_size)])
370            .build()
371            .unwrap();
372
373        let pushed_predicate = PredicateBuilder::new(table.schema().fields())
374            .greater_or_equal("value", Datum::Int(10))
375            .unwrap();
376
377        let schema = Arc::new(ArrowSchema::new(vec![Field::new(
378            "id",
379            ArrowDataType::Int32,
380            false,
381        )]));
382        let scan = PaimonTableScan::new(
383            schema,
384            table,
385            Some(vec!["id".to_string()]),
386            Some(pushed_predicate),
387            vec![Arc::from(vec![split])],
388            None,
389            false,
390        );
391
392        let ctx = SessionContext::new();
393        let stream = scan
394            .execute(0, ctx.task_ctx())
395            .expect("execute should succeed");
396        let batches = stream.try_collect::<Vec<_>>().await.unwrap();
397
398        let actual_ids: Vec<i32> = batches
399            .iter()
400            .flat_map(|batch| {
401                let ids = batch
402                    .column(0)
403                    .as_any()
404                    .downcast_ref::<Int32Array>()
405                    .expect("id column should be Int32Array");
406                (0..ids.len()).map(|idx| ids.value(idx)).collect::<Vec<_>>()
407            })
408            .collect();
409
410        assert_eq!(actual_ids, vec![2, 3, 4]);
411    }
412}
paimon_datafusion/physical_plan/scan.rs

paimon_datafusion/physical_plan/
scan.rs
