kosiew commented on code in PR #26458:
URL: https://github.com/apache/datafusion/pull/26458#discussion_r4312044910
##########
datafusion/datasource/src/file_scan_config/mod.rs:
##########
@@ -836,6 +836,23 @@ impl DataSource for FileScanConfig {
return Ok(None);
}
+ // An ordered scan can use the size-balanced split of an unordered
+ // scan when every new group is still sorted. That split keeps the
+ // files in group order and only moves the group boundaries, so groups
+ // like `[f1, f2], [f3, f4]` with a large `f1` become
+ // `[f1 head], [f1 tail, f2], ...`.
+ if output_ordering.is_some()
+ && let Some(balanced) = self.file_source.repartitioned(
+ target_partitions,
+ repartition_file_min_size,
+ None,
+ self,
+ )?
+ && balanced.validated_output_ordering() ==
self.validated_output_ordering()
Review Comment:
The ordering check can accept an invalid merge when sort keys contain NULLs,
because the min/max validation ignores null counts. For example, with ASC NULLS
LAST, combining `[1, NULL]` and `[2, 3]` can place NULL before non-null values
while still passing validation. Please add a null-aware check that falls back
when ordering cannot be proven, including unknown null counts, without
rejecting valid singleton groups. Also add regression tests for both NULLS
FIRST and NULLS LAST, including a SQL test that reads rows from both files and
verifies the results.
##########
datafusion/datasource/src/file_scan_config/mod.rs:
##########
@@ -3260,6 +3277,89 @@ mod tests {
))
}
+ fn ordered_scan_with_groups(file_groups: Vec<FileGroup>) -> FileScanConfig
{
+ let schema =
+ Arc::new(Schema::new(vec![Field::new("a", DataType::Float64,
false)]));
+ let sort_expr =
PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)));
+ FileScanConfigBuilder::new(
+ ObjectStoreUrl::local_filesystem(),
+ Arc::new(MockSource::new(schema)),
+ )
+ .with_file_groups(file_groups)
+ .with_output_ordering(vec![LexOrdering::new(vec![sort_expr]).unwrap()])
+ .build()
+ }
+
+ fn sized_file_with_stats(
+ name: &str,
+ size: u64,
+ min: f64,
+ max: f64,
+ ) -> PartitionedFile {
+ let mut file = make_file_with_stats(name, min, max);
+ file.object_meta.size = size;
+ file
+ }
+
+ fn group_sizes(config: &FileScanConfig) -> Vec<u64> {
+ config
+ .file_groups
+ .iter()
+ .map(|group| group.iter().map(|f| f.effective_size()).sum())
+ .collect()
+ }
+
+ /// An ordered scan whose files are sorted across groups gets the same
+ /// size-balanced split as an unordered scan and keeps its ordering.
+ #[test]
+ fn repartition_balances_ordered_groups() -> Result<()> {
+ let config = ordered_scan_with_groups(vec![
+ FileGroup::new(vec![
+ sized_file_with_stats("f1", 8000, 0.0, 9.0),
+ sized_file_with_stats("f2", 8000, 10.0, 19.0),
+ ]),
+ FileGroup::new(vec![
+ sized_file_with_stats("f3", 1000, 20.0, 29.0),
+ sized_file_with_stats("f4", 1000, 30.0, 39.0),
+ ]),
+ ]);
+ let ordering = config.output_ordering.first().cloned();
+ assert_eq!(group_sizes(&config), vec![16000, 2000]);
+
+ let repartitioned = DataSource::repartitioned(&config, 4, 0, ordering)?
+ .expect("ordered groups should be rebalanced");
+ let repartitioned = repartitioned
+ .downcast_ref::<FileScanConfig>()
+ .expect("expected FileScanConfig");
+ assert_eq!(group_sizes(repartitioned), vec![4500, 4500, 4500, 4500]);
+ assert_eq!(
+ repartitioned.validated_output_ordering(),
+ config.validated_output_ordering()
+ );
+ Ok(())
+ }
+
+ /// If moving the group boundaries would put overlapping files in one
+ /// group, the scan keeps its groups and its ordering.
+ #[test]
+ fn repartition_keeps_ordered_groups_when_rebalancing_breaks_order() ->
Result<()> {
+ let config = ordered_scan_with_groups(vec![
+ FileGroup::new(vec![
+ sized_file_with_stats("f1", 8000, 0.0, 9.0),
+ sized_file_with_stats("f2", 8000, 10.0, 19.0),
+ ]),
+ FileGroup::new(vec![
+ sized_file_with_stats("f3", 1000, 5.0, 15.0),
+ sized_file_with_stats("f4", 1000, 16.0, 25.0),
+ ]),
+ ]);
+ let ordering = config.output_ordering.first().cloned();
+ assert_eq!(config.validated_output_ordering().len(), 1);
+
+ assert!(DataSource::repartitioned(&config, 4, 0, ordering)?.is_none());
Review Comment:
The current rejection test only covers a case where both repartitioning
strategies return None, so it doesn't verify that the ordered fallback can
succeed. Could you add a mixed-group test using `[[f1(8000, 0..9), f2(8000,
10..19)], [f3(1000, 5..15)]]` with four target partitions, and verify that the
balanced candidate is rejected while the ordered fallback returns Some with the
first group intact, f3 split, and ordering preserved?
##########
datafusion/sqllogictest/test_files/repartition_scan.slt:
##########
Review Comment:
The nearby comment says EXPLAIN should never contain multi-file groups, but
the updated plan intentionally includes one. Could you update it to explain
that groups may span files as long as their concatenation preserves the
declared ordering?
--
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]
