jayzhan211 commented on code in PR #24963:
URL: https://github.com/apache/datafusion/pull/24963#discussion_r3943259941
##########
datafusion/core/tests/physical_optimizer/enforce_distribution.rs:
##########
@@ -2856,6 +2860,142 @@ fn
existing_interleave_is_kept_when_children_stay_interleavable() -> Result<()>
Ok(())
}
+#[test]
+fn interleave_broken_by_later_rewrite_is_repaired_on_next_pass() -> Result<()>
{
+ // The chain reported in https://github.com/apache/datafusion/issues/21826:
+ // a distribution pass builds an interleave, a later rule changes one
+ // child's partitioning while rebuilding the tree, and another distribution
+ // pass runs afterwards.
+ let alias = vec![("a".to_string(), "a1".to_string())];
+ let union = union_exec(vec![
+ aggregate_exec_with_alias(parquet_exec(), alias.clone()),
+ aggregate_exec_with_alias(parquet_exec(), alias),
+ ]);
+ let config = TestConfig::default().config;
+ let pass1 = EnsureRequirements::new().optimize(union, &config)?;
+ assert!(pass1.is::<InterleaveExec>());
+
+ // Stand-in for the later rewrite: one child loses its hash partitioning.
+ // Rebuilding the interleave over it used to fail here.
+ let children = pass1.children();
+ let coalesced: Arc<dyn ExecutionPlan> =
+ Arc::new(CoalescePartitionsExec::new(Arc::clone(children[1])));
+ let rewritten = Arc::clone(&pass1).replace_children(
+ vec![Arc::clone(children[0]), coalesced],
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+ )?;
+ assert!(matches!(
+ rewritten.output_partitioning(),
+ Partitioning::UnknownPartitioning(10)
+ ));
+ // Without a repair, the plan is rejected rather than executed.
+ assert!(
+ SanityCheckPlan::new()
+ .optimize(Arc::clone(&rewritten), &config)
+ .is_err()
+ );
+
+ // The next distribution pass restores a valid interleave.
+ let pass2 = EnsureRequirements::new().optimize(rewritten, &config)?;
+ SanityCheckPlan::new().optimize(Arc::clone(&pass2), &config)?;
+ assert_plan!(pass2,
+ @r"
+ InterleaveExec
+ AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[]
+ RepartitionExec: partitioning=Hash([a1@0], 10), input_partitions=10
+ AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[]
+ RepartitionExec: partitioning=RoundRobinBatch(10),
input_partitions=1
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b,
c, d, e], file_type=parquet
+ AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[]
+ RepartitionExec: partitioning=Hash([a1@0], 10), input_partitions=10
+ AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[]
+ RepartitionExec: partitioning=RoundRobinBatch(10),
input_partitions=1
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b,
c, d, e], file_type=parquet
+ ");
+
+ Ok(())
+}
+
+/// A single-partition parquet scan with the given (inexact) row and byte
+/// statistics, so `JoinSelection` can compare build and probe sizes.
+fn parquet_exec_with_size(
+ num_rows: usize,
+ total_byte_size: usize,
+) -> Arc<dyn ExecutionPlan> {
+ let mut statistics = Statistics::new_unknown(&schema());
+ statistics.num_rows = Precision::Inexact(num_rows);
+ statistics.total_byte_size = Precision::Inexact(total_byte_size);
+ let config = FileScanConfigBuilder::new(
+ ObjectStoreUrl::parse("test:///").unwrap(),
+ Arc::new(ParquetSource::new(schema())),
+ )
+ .with_file(PartitionedFile::new(
+ "x".to_string(),
+ total_byte_size as u64,
+ ))
+ .with_statistics(statistics)
+ .build();
+ DataSourceExec::from_data_source(config)
+}
+
+#[test]
+fn issue_21826_join_selection_after_distribution_pass() -> Result<()> {
Review Comment:
The case in #21826
--
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]
