mbutrovich commented on code in PR #55: URL: https://github.com/apache/datafusion-iceberg/pull/55#discussion_r4298479904
########## 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: What do you think about adding a runnable example here? Users interact with this type in two steps: register it on the session, then `SET` an option. An example that does both and reads the value back shows the registration step the prose describes. It also runs as a doctest, so it can't go stale. ```suggestion /// such as `max_merge_files`, and `SHOW` and `RESET` do not support them. /// /// ``` /// # #[tokio::main(flavor = "current_thread")] /// # async fn main() -> datafusion::error::Result<()> { /// use datafusion::prelude::{SessionConfig, SessionContext}; /// use iceberg_datafusion::IcebergDataFusionConfig; /// /// let config = /// SessionConfig::new().with_option_extension(IcebergDataFusionConfig::default()); /// let ctx = SessionContext::new_with_config(config); /// ctx.sql("SET iceberg.planning.preserve_data_ordering = true") /// .await? /// .collect() /// .await?; /// /// let state = ctx.state(); /// let options = state.config().options(); /// let iceberg = options.extensions.get::<IcebergDataFusionConfig>().unwrap(); /// assert!(iceberg.planning.preserve_data_ordering); /// # Ok(()) /// # } /// ``` ``` ########## 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: Should the group follow DataFusion's layout? DataFusion puts its options under fixed groups ([`ConfigOptions`](https://github.com/apache/datafusion/blob/d5552342012888b7d1a3ab88d92e3d292fc0cde0/datafusion/common/src/config.rs#L1915-L1931): `catalog`, `execution`, `optimizer`, `sql_parser`, `explain`, `format`), and none of them is `planning`. That name comes from Spark's `spark.sql.iceberg.planning.*`. The DataFusion option closest to this one is [`split_file_groups_by_statistics`](https://github.com/apache/datafusion/blob/d5552342012888b7d1a3ab88d92e3d292fc0cde0/datafusion/common/src/config.rs#L991-L994), which also arranges files so that a sort can be removed. It lives under `execution`, next to the Parquet reader options in `execution.parquet`. This crate may later expose iceberg-rust's reader settings as session options ([`ArrowReaderBuilder`](https://github.com/apache/iceberg-rust/blob/5b82ab99aeaf50bf4f5b3b776ac9cd6380a5f3e3/crates/iceberg/src/arrow/reader/mod.rs#L83-L117): `data_file_concurrency_limit`, `row_group_filtering_enabled`, and so on). If so, `iceberg.execution.*` with iceberg-rust's names could hold those too. What do you think about `iceberg.execution.preserve_data_ordering`? ########## 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: Should this say where the nulls-first default comes from? The spec has no default null order. A sort field's [`null-order`](https://github.com/apache/iceberg/blob/524fe9ef1b4aa72be19623f667738b21a02b23b2/format/spec.md?plain=1#L1941) is required, and the default comes from writers such as Iceberg Java's [`SortOrderBuilder.asc`](https://github.com/apache/iceberg/blob/524fe9ef1b4aa72be19623f667738b21a02b23b2/api/src/main/java/org/apache/iceberg/SortOrderBuilder.java#L33-L35). For a table whose order is `id ASC NULLS LAST`, a plain `ORDER BY id` can skip its sort, which this paragraph currently rules out. The PR description has the same sentence. ```suggestion /// The scan reports each sort field's null order as the table declares /// it. Writers commonly sort ascending fields with nulls first, the /// default of Iceberg Java's `SortOrderBuilder.asc`, while DataFusion's /// `ORDER BY c` puts them last. On a nullable column sorted that way, /// 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. ``` ########## 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: Could this be named for what it caps once the scan has more than one partition? The PR description says that once apache/iceberg-rust#2671 groups files into partitions, each group can be merged the same way, and from then on this cap applies per partition. apache/datafusion-comet#5331 names the same limit, with the same default of 64, [`sortMerge.maxFilesPerPartition`](https://github.com/apache/datafusion-comet/blob/4a45cd6917ab79374955bb03eea3eed5e483785d/spark/src/main/scala/org/apache/comet/CometConf.scala#L138-L139). Spark's in apache/iceberg#16750 is [`preserve-data-ordering.max-files-per-task`](https://github.com/apache/iceberg/blob/9223faa7523de6884309ecfc08c8a6a00de5b817/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java#L49-L51). `max_files_per_partition` would match Comet and DataFusion's own use of "partition", and it would stay accurate after apache/iceberg-rust#2671 without a rename. ########## 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: Should this compare the sort orders themselves rather than their ids? `FileScanTask` carries both the raw `sort_order_id` and a resolved `sort_order`, and its [validation](https://github.com/apache/iceberg-rust/blob/5b82ab99aeaf50bf4f5b3b776ac9cd6380a5f3e3/crates/iceberg/src/scan/task.rs#L271-L296) doesn't check that they agree. Tasks planned from one table always agree. Tasks passed to `IcebergDataSource::builder` come from outside, though, and `build()` documents that it fails if they "don't share a sort order the source can report". As written, two tasks with the same id but different orders pass, and the merge uses the first task's order for both. Comparing the orders checks what the merge actually relies on, and needs one condition instead of two. ```suggestion if !tasks.iter().all(|task| task.sort_order() == Some(order)) { ``` ########## 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: What do you think about adding tests for a sorted scan of one file and of no files? Every test that reports an order merges at least two files, so nothing exercises this early return. To check, I made it return an empty stream (`EmptyRecordBatchStream::new(self.schema())`). The integration tests and all the sqllogictest schedules still pass. `create_sorted_table` with a single `(SORTED, &[...])` entry, and with `&[]`, would cover both cases. One file should report the order and return its rows. An empty table should report no order and return no rows. -- 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]
