jayzhan211 commented on code in PR #26458:
URL: https://github.com/apache/datafusion/pull/26458#discussion_r4336728739
##########
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:
Added
`repartition_falls_back_to_ordered_split_when_rebalancing_breaks_order` with
this layout and 4 target partitions. The balanced split is rejected, and the
ordered split keeps `[f1, f2]` as the first group, splits `f3` into the other
three groups, and keeps the 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]
