anuragmantri commented on code in PR #55: URL: https://github.com/apache/datafusion-iceberg/pull/55#discussion_r4299373466
########## crates/datafusion/src/config.rs: ########## @@ -0,0 +1,66 @@ +// 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. + +//! Session options for Iceberg tables, set with `SET iceberg.<group>.<option>`. + +use datafusion::common::config::ConfigExtension; +use datafusion::common::extensions_options; + +extensions_options! { + /// Session options for reading Iceberg tables through DataFusion. + /// + /// `SET iceberg.<group>.<option> = <value>` only works once this extension + /// is registered on the session, with + /// [`SessionConfig::with_option_extension`](datafusion::prelude::SessionConfig::with_option_extension). + /// Without it, every option keeps its default. DataFusion lists these + /// options in `information_schema.df_settings` by their last name alone, + /// such as `max_merge_files`, and `SHOW` and `RESET` do not support them. Review Comment: Added your example in 6a0c8f1, with the new option names. It runs as a doctest and passes. ########## AGENTS.md: ########## @@ -73,6 +73,8 @@ Tests are self-contained (in-memory catalogs, temp dirs, checked-in metadata JSO `IcebergTableScan` converts the projection to column names and the filters to a single Iceberg `Predicate` (`physical_plan/expr_to_predicate.rs`). It applies `limit` in-stream, has one output partition, and delegates the actual read to iceberg's `TableScan::to_arrow()`. +With the session option `iceberg.planning.preserve_data_ordering` (`config.rs`; off by default), `TableProvider::scan` lists the data files up front. If there are at most `iceberg.planning.max_merge_files` of them and all record the same resolvable sort order, `IcebergDataSource` reports that order's leading fields as its output ordering and k-way merges one reader per file in `open()`. Only identity transforms of projected top-level columns are reported, stopping at the first float, double, or UUID field, because Arrow orders -0.0 and NaN differently from Spark-written files. Otherwise the source is built exactly as without the option. A source rebuilt with `IcebergDataSource::builder` must be given what `sorted_tasks()` returns through `with_sorted_tasks`, or it will read the files unordered under a plan that relies on their order. The order is also not reported unless the provider's declared schema matches the snapshot schema the files are read with, since DataFusion ignores the null order of a column declared non-nullable. Review Comment: Done in 6370be9, thanks. ########## crates/datafusion/src/config.rs: ########## @@ -0,0 +1,66 @@ +// 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. + +//! Session options for Iceberg tables, set with `SET iceberg.<group>.<option>`. + +use datafusion::common::config::ConfigExtension; +use datafusion::common::extensions_options; + +extensions_options! { + /// Session options for reading Iceberg tables through DataFusion. + /// + /// `SET iceberg.<group>.<option> = <value>` only works once this extension + /// is registered on the session, with + /// [`SessionConfig::with_option_extension`](datafusion::prelude::SessionConfig::with_option_extension). + /// Without it, every option keeps its default. DataFusion lists these + /// options in `information_schema.df_settings` by their last name alone, + /// such as `max_merge_files`, and `SHOW` and `RESET` do not support them. + pub struct IcebergDataFusionConfig { + /// Options for planning table scans. + pub planning: IcebergPlanningConfig, default = IcebergPlanningConfig::default() + } +} + +extensions_options! { + /// Options for planning Iceberg table scans, under `iceberg.planning`. + pub struct IcebergPlanningConfig { + /// When true, a scan lists its data files while it is planned and, if + /// every file records the same sort order, reports that order to + /// DataFusion so that it can skip sorting the scan's output. The scan + /// then merges the files in that order, with every file open at once. + /// + /// Only the leading sort fields that are identity transforms of + /// projected top-level columns are reported, stopping at the first + /// floating point or UUID field, whose orders in Iceberg writers and in + /// DataFusion can differ. + /// + /// Iceberg sorts ascending fields with nulls first by default, while + /// DataFusion's `ORDER BY c` puts them last. On a nullable sort column, + /// only `ORDER BY c NULLS FIRST`, or a session with + /// `datafusion.sql_parser.default_null_ordering = 'nulls_min'`, can + /// skip its sort. Other queries sort the merged output again. Review Comment: You're right, the spec requires a null order on every sort field and has no default. Applied your wording in 6a0c8f1, and I'll fix the same sentence in the description. ########## crates/datafusion/src/config.rs: ########## @@ -0,0 +1,66 @@ +// 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. + +//! Session options for Iceberg tables, set with `SET iceberg.<group>.<option>`. + +use datafusion::common::config::ConfigExtension; +use datafusion::common::extensions_options; + +extensions_options! { + /// Session options for reading Iceberg tables through DataFusion. + /// + /// `SET iceberg.<group>.<option> = <value>` only works once this extension + /// is registered on the session, with + /// [`SessionConfig::with_option_extension`](datafusion::prelude::SessionConfig::with_option_extension). + /// Without it, every option keeps its default. DataFusion lists these + /// options in `information_schema.df_settings` by their last name alone, + /// such as `max_merge_files`, and `SHOW` and `RESET` do not support them. + pub struct IcebergDataFusionConfig { + /// Options for planning table scans. + pub planning: IcebergPlanningConfig, default = IcebergPlanningConfig::default() + } +} + +extensions_options! { + /// Options for planning Iceberg table scans, under `iceberg.planning`. + pub struct IcebergPlanningConfig { + /// When true, a scan lists its data files while it is planned and, if + /// every file records the same sort order, reports that order to + /// DataFusion so that it can skip sorting the scan's output. The scan + /// then merges the files in that order, with every file open at once. + /// + /// Only the leading sort fields that are identity transforms of + /// projected top-level columns are reported, stopping at the first + /// floating point or UUID field, whose orders in Iceberg writers and in + /// DataFusion can differ. + /// + /// Iceberg sorts ascending fields with nulls first by default, while + /// DataFusion's `ORDER BY c` puts them last. On a nullable sort column, + /// only `ORDER BY c NULLS FIRST`, or a session with + /// `datafusion.sql_parser.default_null_ordering = 'nulls_min'`, can + /// skip its sort. Other queries sort the merged output again. + pub preserve_data_ordering: bool, default = false + /// The most data files a scan merges to preserve their sort order. + /// A scan of more files reports no order, and DataFusion sorts its + /// output instead, which can spill to disk where the merge cannot. + pub max_merge_files: usize, default = 64 Review Comment: Agreed on naming it per partition, so it stays accurate after apache/iceberg-rust#2671. I kept "merge" in it, though: `max_merge_files_per_partition`. Under `iceberg.execution`, a bare `max_files_per_partition` reads like a general limit on how files are grouped into partitions, which is the kind of option apache/iceberg-rust#2671 might add itself. Happy to drop "merge" if you prefer Comet's exact name. ########## crates/datafusion/src/read/scan.rs: ########## @@ -93,13 +134,124 @@ 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. + pub fn sorted_tasks(&self) -> Option<&[FileScanTask]> { + self.sorted.as_ref().map(|sorted| sorted.tasks.as_slice()) + } + + /// Lists the source's data files and reports the sort order they all + /// record, if the session enables `iceberg.planning.preserve_data_ordering` + /// and they are few enough to merge. Otherwise returns the source unchanged. + pub(crate) async fn plan_sort_order(self, state: &dyn Session) -> Result<Self> { + let config = state + .config() + .options() + .extensions + .get::<IcebergDataFusionConfig>() + .map(|config| config.planning.clone()) + .unwrap_or_default(); + if !config.preserve_data_ordering { + return Ok(self); + } + + let metadata = self.table.metadata(); + let snapshot = match self.snapshot_id { + Some(snapshot_id) => metadata.snapshot_by_id(snapshot_id), + None => metadata.current_snapshot(), + }; + let Some(snapshot) = snapshot else { + return Ok(self); + }; + let schema = snapshot.schema(metadata).map_err(to_datafusion_error)?; + let output = self.schema(); + if !metadata + .sort_orders_iter() + .any(|order| lex_ordering(order, &schema, &output).is_some()) + { + return Ok(self); + } + + let tasks: Vec<FileScanTask> = build_table_scan( + &self.table, + self.snapshot_id, + Some(self.projection.clone()), + self.predicate.clone(), + )? + .plan_files() + .await + .map_err(to_datafusion_error)? + .take(config.max_merge_files.saturating_add(1)) + .try_collect() + .await + .map_err(to_datafusion_error)?; + if tasks.len() > config.max_merge_files { + return Ok(self); + } + let mut source = self; + source.sorted = shared_ordering(&tasks, &output) + .map(|ordering| SortedTasks { tasks, ordering }); + Ok(source) + } + + /// 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> { + // The builder defaults to the reader settings `TableScan::to_arrow` uses. + let reader = self.table.reader_builder().build(); + let mut streams = sorted + .tasks + .iter() + .map(|task| { + let tasks = futures::stream::iter([Ok(task.clone())]).boxed(); + let batches = reader + .clone() + .read(tasks) + .map_err(to_datafusion_error)? + .stream() + .map_err(to_datafusion_error); + Ok( + Box::pin(RecordBatchStreamAdapter::new(self.schema(), batches)) + as SendableRecordBatchStream, + ) + }) + .collect::<Result<Vec<_>>>()?; + if streams.len() == 1 { + return Ok(streams.remove(0)); + } Review Comment: Good catch. Added `test_sorted_scan_of_one_file_and_of_no_files` in 6a0c8f1. One file reports the order and returns its rows, and an empty table reports no order and returns none. With your mutation (an empty stream for one file) that test fails. ########## crates/datafusion/src/read/scan.rs: ########## @@ -287,12 +500,300 @@ async fn get_batch_stream( if let Some(predicate) = predicate { scan_builder = scan_builder.with_filter(predicate); } - let table_scan = scan_builder.build().map_err(to_datafusion_error)?; + scan_builder.build().map_err(to_datafusion_error) +} - let stream = table_scan - .to_arrow() - .await - .map_err(to_datafusion_error)? - .map_err(to_datafusion_error); - Ok(Box::pin(stream)) +/// The order of the columns of `output` that every one of `tasks` is sorted +/// in, or `None` if they record different sort orders, one records none, or +/// the schema they are read with does not match `output`. +fn shared_ordering(tasks: &[FileScanTask], output: &ArrowSchema) -> Option<LexOrdering> { + let first = tasks.first()?; + let order = first.sort_order()?; + if !tasks.iter().all(|task| { + task.sort_order_id() == first.sort_order_id() && task.sort_order().is_some() + }) { Review Comment: Agreed, applied in 6a0c8f1. Added `test_rebuilt_sorted_scan_rejects_tasks_with_different_orders`, which gives two tasks the same sort order id but different orders and expects `build()` to fail. Going back to comparing ids fails it. ########## crates/datafusion/src/config.rs: ########## @@ -0,0 +1,66 @@ +// 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. + +//! Session options for Iceberg tables, set with `SET iceberg.<group>.<option>`. + +use datafusion::common::config::ConfigExtension; +use datafusion::common::extensions_options; + +extensions_options! { + /// Session options for reading Iceberg tables through DataFusion. + /// + /// `SET iceberg.<group>.<option> = <value>` only works once this extension + /// is registered on the session, with + /// [`SessionConfig::with_option_extension`](datafusion::prelude::SessionConfig::with_option_extension). + /// Without it, every option keeps its default. DataFusion lists these + /// options in `information_schema.df_settings` by their last name alone, + /// such as `max_merge_files`, and `SHOW` and `RESET` do not support them. + pub struct IcebergDataFusionConfig { + /// Options for planning table scans. + pub planning: IcebergPlanningConfig, default = IcebergPlanningConfig::default() Review Comment: Agreed, the group is now `execution` in 6a0c8f1, so the options are `iceberg.execution.preserve_data_ordering` and `iceberg.execution.max_merge_files_per_partition`. The struct is renamed `IcebergExecutionConfig`. Following DataFusion's layout makes more sense here than following Spark's, and it leaves room for iceberg-rust's reader settings. -- 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]
