iceberg_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::pin::Pin;
19use std::sync::Arc;
20use std::vec;
21
22use datafusion::arrow::array::RecordBatch;
23use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef;
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 datafusion::prelude::Expr;
31use futures::{Stream, TryStreamExt};
32use iceberg::expr::Predicate;
33use iceberg::table::Table;
34
35use super::expr_to_predicate::convert_filters_to_predicate;
36use crate::to_datafusion_error;
37
38/// Manages the scanning process of an Iceberg [`Table`], encapsulating the
39/// necessary details and computed properties required for execution planning.
40#[derive(Debug)]
41pub struct IcebergTableScan {
42    /// A table in the catalog.
43    table: Table,
44    /// Snapshot of the table to scan.
45    snapshot_id: Option<i64>,
46    /// Stores certain, often expensive to compute,
47    /// plan properties used in query optimization.
48    plan_properties: Arc<PlanProperties>,
49    /// Projection column names, None means all columns
50    projection: Option<Vec<String>>,
51    /// Filters to apply to the table scan
52    predicates: Option<Predicate>,
53    /// Optional limit on the number of rows to return
54    limit: Option<usize>,
55}
56
57impl IcebergTableScan {
58    /// Creates a new [`IcebergTableScan`] object.
59    pub(crate) fn new(
60        table: Table,
61        snapshot_id: Option<i64>,
62        schema: ArrowSchemaRef,
63        projection: Option<&Vec<usize>>,
64        filters: &[Expr],
65        limit: Option<usize>,
66    ) -> Self {
67        let output_schema = match projection {
68            None => schema.clone(),
69            Some(projection) => Arc::new(schema.project(projection).unwrap()),
70        };
71        let plan_properties = Self::compute_properties(output_schema.clone());
72        let projection = get_column_names(schema.clone(), projection);
73        let predicates = convert_filters_to_predicate(filters);
74
75        Self {
76            table,
77            snapshot_id,
78            plan_properties,
79            projection,
80            predicates,
81            limit,
82        }
83    }
84
85    pub fn table(&self) -> &Table {
86        &self.table
87    }
88
89    pub fn snapshot_id(&self) -> Option<i64> {
90        self.snapshot_id
91    }
92
93    pub fn projection(&self) -> Option<&[String]> {
94        self.projection.as_deref()
95    }
96
97    pub fn predicates(&self) -> Option<&Predicate> {
98        self.predicates.as_ref()
99    }
100
101    pub fn limit(&self) -> Option<usize> {
102        self.limit
103    }
104
105    /// Computes [`PlanProperties`] used in query optimization.
106    fn compute_properties(schema: ArrowSchemaRef) -> Arc<PlanProperties> {
107        // TODO:
108        // This is more or less a placeholder, to be replaced
109        // once we support output-partitioning
110        Arc::new(PlanProperties::new(
111            EquivalenceProperties::new(schema),
112            Partitioning::UnknownPartitioning(1),
113            EmissionType::Incremental,
114            Boundedness::Bounded,
115        ))
116    }
117}
118
119impl ExecutionPlan for IcebergTableScan {
120    fn name(&self) -> &str {
121        "IcebergTableScan"
122    }
123
124    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan + 'static>> {
125        vec![]
126    }
127
128    fn with_new_children(
129        self: Arc<Self>,
130        _children: Vec<Arc<dyn ExecutionPlan>>,
131    ) -> DFResult<Arc<dyn ExecutionPlan>> {
132        Ok(self)
133    }
134
135    fn properties(&self) -> &Arc<PlanProperties> {
136        &self.plan_properties
137    }
138
139    fn execute(
140        &self,
141        _partition: usize,
142        _context: Arc<TaskContext>,
143    ) -> DFResult<SendableRecordBatchStream> {
144        let fut = get_batch_stream(
145            self.table.clone(),
146            self.snapshot_id,
147            self.projection.clone(),
148            self.predicates.clone(),
149        );
150        let stream = futures::stream::once(fut).try_flatten();
151
152        // Apply limit if specified
153        let limited_stream: Pin<Box<dyn Stream<Item = DFResult<RecordBatch>> + Send>> =
154            if let Some(limit) = self.limit {
155                let mut remaining = limit;
156                Box::pin(stream.try_filter_map(move |batch| {
157                    futures::future::ready(if remaining == 0 {
158                        Ok(None)
159                    } else if batch.num_rows() <= remaining {
160                        remaining -= batch.num_rows();
161                        Ok(Some(batch))
162                    } else {
163                        let limited_batch = batch.slice(0, remaining);
164                        remaining = 0;
165                        Ok(Some(limited_batch))
166                    })
167                }))
168            } else {
169                Box::pin(stream)
170            };
171
172        Ok(Box::pin(RecordBatchStreamAdapter::new(
173            self.schema(),
174            limited_stream,
175        )))
176    }
177}
178
179impl DisplayAs for IcebergTableScan {
180    fn fmt_as(
181        &self,
182        _t: datafusion::physical_plan::DisplayFormatType,
183        f: &mut std::fmt::Formatter,
184    ) -> std::fmt::Result {
185        write!(
186            f,
187            "IcebergTableScan projection:[{}] predicate:[{}]",
188            self.projection
189                .clone()
190                .map_or(String::new(), |v| v.join(",")),
191            self.predicates
192                .clone()
193                .map_or(String::from(""), |p| format!("{p}"))
194        )?;
195        if let Some(limit) = self.limit {
196            write!(f, " limit:[{limit}]")?;
197        }
198        Ok(())
199    }
200}
201
202/// Asynchronously retrieves a stream of [`RecordBatch`] instances
203/// from a given table.
204///
205/// This function initializes a [`TableScan`], builds it,
206/// and then converts it into a stream of Arrow [`RecordBatch`]es.
207async fn get_batch_stream(
208    table: Table,
209    snapshot_id: Option<i64>,
210    column_names: Option<Vec<String>>,
211    predicates: Option<Predicate>,
212) -> DFResult<Pin<Box<dyn Stream<Item = DFResult<RecordBatch>> + Send>>> {
213    let scan_builder = match snapshot_id {
214        Some(snapshot_id) => table.scan().snapshot_id(snapshot_id),
215        None => table.scan(),
216    };
217
218    let mut scan_builder = match column_names {
219        Some(column_names) => scan_builder.select(column_names),
220        None => scan_builder.select_all(),
221    };
222    if let Some(pred) = predicates {
223        scan_builder = scan_builder.with_filter(pred);
224    }
225    let table_scan = scan_builder.build().map_err(to_datafusion_error)?;
226
227    let stream = table_scan
228        .to_arrow()
229        .await
230        .map_err(to_datafusion_error)?
231        .map_err(to_datafusion_error);
232    Ok(Box::pin(stream))
233}
234
235fn get_column_names(
236    schema: ArrowSchemaRef,
237    projection: Option<&Vec<usize>>,
238) -> Option<Vec<String>> {
239    projection.map(|v| {
240        v.iter()
241            .map(|p| schema.field(*p).name().clone())
242            .collect::<Vec<String>>()
243    })
244}
iceberg_datafusion/physical_plan/scan.rs

iceberg_datafusion/physical_plan/
scan.rs
