gabotechs commented on code in PR #55:
URL: https://github.com/apache/datafusion-iceberg/pull/55#discussion_r4322184389
##########
crates/datafusion/src/table_provider/catalog_table_provider.rs:
##########
@@ -139,13 +140,12 @@ impl TableProvider for IcebergCatalogTableProvider {
.map_err(to_datafusion_error)?;
// Create scan with fresh metadata (always use current snapshot)
- let source = IcebergDataSource::builder(table, self.schema.clone(),
None)
+ let builder = IcebergDataSource::builder(table, self.schema.clone())
.with_projection(projection)
.with_filters(filters)
- .with_fetch(limit)
- .build()?
- .plan_sort_order(state)
- .await?;
+ .with_fetch(limit);
+ let sorted_tasks = plan_sorted_tasks(&builder, state).await?;
+ let source = builder.with_sorted_tasks(sorted_tasks).build()?;
Review Comment:
There's still an issue with the API here: `builder.with_sorted_tasks` is
unusable from the outside, because it requires someone to call
`plan_sorted_tasks` with is `pub(crate)`.
An clean builder-pattern API should look something like this:
```rust
let source = IcebergDataSource::builder(self.table.clone(),
self.schema.clone())
.with_snapshot_id(self.snapshot_id)
.with_projection(projection)
.with_filters(filters)
.with_fetch(limit)
.with_materialized_sorted_tasks(state).await?
.build();
```
Without requiring users from calling private methods to use the API.
##########
crates/datafusion/src/read/scan.rs:
##########
@@ -93,14 +121,73 @@ impl IcebergDataSource {
pub fn fetch(&self) -> Option<usize> {
self.fetch
}
+
+ /// The data files the source merges in the order it reports, listed while
+ /// the scan was planned, or `None` if it reports no order and lists its
+ /// files when it runs. A source rebuilt from this one must be given them
+ /// with [`IcebergDataSourceBuilder::with_sorted_tasks`], because a plan
+ /// built over a sorted source may rely on its order.
+ pub fn sorted_tasks(&self) -> Option<&[FileScanTask]> {
+ self.sorted.as_ref().map(|sorted| sorted.tasks.as_slice())
+ }
+
+ /// Reads each sorted data file as its own stream and merges the streams in
+ /// the order the source reports.
+ fn merge_sorted(
+ &self,
+ sorted: &SortedTasks,
+ partition: usize,
+ context: &TaskContext,
+ ) -> Result<SendableRecordBatchStream> {
Review Comment:
As soon as the DataSource starts supporting some partitioning greater than
1, the current code here should break, as we'd be returning duplicate data in
each partition.
##########
crates/datafusion/tests/sorted_iceberg_scan.rs:
##########
@@ -0,0 +1,1002 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Integration tests for sorted scans: with
+//! `iceberg.execution.preserve_data_ordering`, a scan of files that all record
+//! the same sort order reports that order to DataFusion and merges the files.
+
+use std::collections::HashMap;
+use std::error::Error;
+use std::str::FromStr;
+use std::sync::Arc;
+
+use datafusion::arrow::array::{
+ AsArray, Int32Array, Int64Array, ListArray, RecordBatch, StringArray,
StructArray,
+};
+use datafusion::arrow::buffer::OffsetBuffer;
+use datafusion::arrow::compute::concat_batches;
+use datafusion::arrow::datatypes::{DataType, Field, Int32Type, Schema as
ArrowSchema};
+use datafusion::datasource::TableProvider;
+use datafusion::datasource::source::{DataSource, DataSourceExec};
+use datafusion::execution::context::{SessionConfig, SessionContext};
+use datafusion::parquet::arrow::{ArrowWriter, PARQUET_FIELD_ID_META_KEY};
+use datafusion::physical_plan::common::collect;
+use datafusion::physical_plan::joins::SortMergeJoinExec;
+use datafusion::physical_plan::sorts::sort::SortExec;
+use datafusion::physical_plan::{ExecutionPlan, displayable};
+use futures::{StreamExt, TryStreamExt};
+use iceberg::arrow::schema_to_arrow_schema;
+use iceberg::io::LocalFsStorageFactory;
+use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder};
+use iceberg::scan::{FileScanTask, FileScanTaskDeleteFile};
+use iceberg::spec::{
+ DataContentType, DataFileBuilder, DataFileFormat, ListType, NestedField,
NullOrder,
+ PrimitiveType, Schema, SortDirection, SortField, SortOrder, SortOrderRef,
Struct,
+ StructType, Transform, Type,
+};
+use iceberg::table::Table;
+use iceberg::transaction::{ApplyTransactionAction, Transaction};
+use iceberg::{
+ Catalog, CatalogBuilder, MemoryCatalog, MetadataLocation, NamespaceIdent,
+ Result as IcebergResult, TableCreation,
+};
+use iceberg_datafusion::read::IcebergDataSource;
+use iceberg_datafusion::{
+ IcebergCatalogTableProvider, IcebergDataFusionConfig,
IcebergStaticTableProvider,
+};
+use tempfile::TempDir;
+
Review Comment:
This file is not really in a reviewable state, I see the amount of synthetic
preparation is disproportionate for what's being tested.
My recommendation is to find some inspiration in how other engines
supporting iceberg test this, and try to minimize the amount of preparatory
boilerplate in order to make the integration tests closer to real situations.
One specific suggestion that comes to mind is:
- For the integration tests, focus on real scenarios, with a real dataset
like one you would find on a production setup (in a smaller case of course) and
cover as much as possible there.
- For the unit test, focus on edge case scenarios that are hard to replicate
otherwise, like the ones needing from very specific shapes of the data.
##########
crates/datafusion/src/read/scan.rs:
##########
@@ -93,14 +121,73 @@ impl IcebergDataSource {
pub fn fetch(&self) -> Option<usize> {
self.fetch
}
+
+ /// The data files the source merges in the order it reports, listed while
+ /// the scan was planned, or `None` if it reports no order and lists its
+ /// files when it runs. A source rebuilt from this one must be given them
+ /// with [`IcebergDataSourceBuilder::with_sorted_tasks`], because a plan
+ /// built over a sorted source may rely on its order.
+ pub fn sorted_tasks(&self) -> Option<&[FileScanTask]> {
+ self.sorted.as_ref().map(|sorted| sorted.tasks.as_slice())
+ }
+
+ /// Reads each sorted data file as its own stream and merges the streams in
+ /// the order the source reports.
+ fn merge_sorted(
Review Comment:
I see the execution mode when files are sorted here is completely different,
and this DataSource implementation is filled with conditionals over the sorted
property for executing different code paths.
This does not seem clean, and I foresee this taking a hit in maintainability
long term. I think we should just create a different `DataSource`
implementation dedicated to sorted scans, something like
`IcebergSortedDataSource` in `read/sorted_scan.rs`, and avoid increasing the
complexity of this file.
Hopefully this will provide a greater sense of isolation of responsibilities.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
