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

iceberg_datafusion/physical_plan/
scan.rs
