gstamatakis95 commented on code in PR #25554:
URL: https://github.com/apache/datafusion/pull/25554#discussion_r4370919935
##########
datafusion/physical-plan/src/joins/utils.rs:
##########
@@ -4648,6 +4713,332 @@ mod tests {
Ok(())
}
+ fn int32_schema(first: &str, second: &str) -> Arc<Schema> {
+ Arc::new(Schema::new(vec![
+ Field::new(first, DataType::Int32, true),
+ Field::new(second, DataType::Int32, true),
+ ]))
+ }
+
+ fn unknown_partitioned(
+ schema: Arc<Schema>,
+ partitions: usize,
+ ) -> Arc<dyn ExecutionPlan> {
+ TestMemoryExec::try_new_exec(&vec![vec![]; partitions], schema, None)
+ .expect("memory source")
+ }
+
+ fn repartitioned(
+ schema: Arc<Schema>,
+ partitioning: Partitioning,
+ ) -> Arc<dyn ExecutionPlan> {
+ let source = unknown_partitioned(schema, 1);
+ Arc::new(RepartitionExec::try_new(source,
partitioning).expect("repartition"))
+ }
+
+ fn col_at(name: &str, index: usize) -> PhysicalExprRef {
+ Arc::new(Column::new(name, index))
+ }
+
+ fn coalesced(first: PhysicalExprRef, second: PhysicalExprRef) ->
PhysicalExprRef {
+ coalesce_keys(first, second).expect("coalesce")
+ }
+
+ fn range_on(
+ expr: PhysicalExprRef,
+ options: SortOptions,
+ split_points: &[i32],
+ ) -> Partitioning {
+ let ordering =
+ LexOrdering::new([PhysicalSortExpr::new(expr,
options)]).expect("one key");
+ let split_points = split_points
+ .iter()
+ .map(|value|
SplitPoint::new(vec![ScalarValue::Int32(Some(*value))]))
+ .collect();
+ let range = RangePartitioning::try_new(ordering,
split_points).expect("range");
+ Partitioning::Range(range)
+ }
+
+ fn a_equals_c() -> JoinOn {
+ vec![(col_at("a", 0), col_at("c", 0))]
+ }
+
+ fn join_partitioning(
+ left: &Arc<dyn ExecutionPlan>,
+ right: &Arc<dyn ExecutionPlan>,
+ join_type: JoinType,
+ on: JoinOnRef,
+ ) -> Partitioning {
+ let (schema, _) = build_join_schema(&left.schema(), &right.schema(),
&join_type);
+ let mut eq_properties = EquivalenceProperties::new(Arc::new(schema));
+ symmetric_join_output_partitioning(
+ left,
+ right,
+ &join_type,
+ on,
+ &mut eq_properties,
+ )
+ .expect("join partitioning")
+ }
+
+ fn full_join_partitioning(
+ left: &Arc<dyn ExecutionPlan>,
+ right: &Arc<dyn ExecutionPlan>,
+ on: JoinOnRef,
+ ) -> Partitioning {
+ join_partitioning(left, right, JoinType::Full, on)
+ }
+
+ #[test]
+ fn full_join_output_partitioning_coalesces_hash_keys() {
+ let left = repartitioned(
+ int32_schema("a", "b"),
+ Partitioning::Hash(vec![col_at("a", 0)], 4),
+ );
+ let right = repartitioned(
+ int32_schema("c", "d"),
+ Partitioning::Hash(vec![col_at("c", 0)], 4),
+ );
+ let on = a_equals_c();
+
+ assert_eq!(
+ full_join_partitioning(&left, &right, &on),
+ Partitioning::Hash(vec![coalesced(col_at("a", 0), col_at("c",
2))], 4)
+ );
+ assert_eq!(
+ join_partitioning(&left, &right, JoinType::Inner, &on),
+ Partitioning::Hash(vec![col_at("c", 2)], 4)
+ );
+ assert_eq!(
+ join_partitioning(&left, &right, JoinType::Left, &on),
+ Partitioning::Hash(vec![col_at("a", 0)], 4)
+ );
+ }
+
+ #[test]
+ fn full_join_output_partitioning_coalesces_range_keys() {
+ let options = SortOptions::new(false, true);
+ let left = repartitioned(
+ int32_schema("a", "b"),
+ range_on(col_at("a", 0), options, &[10, 20, 30]),
+ );
+ let right = repartitioned(
+ int32_schema("c", "d"),
+ range_on(col_at("c", 0), options, &[10, 20, 30]),
+ );
+
+ let expected = range_on(
+ coalesced(col_at("a", 0), col_at("c", 2)),
+ options,
+ &[10, 20, 30],
+ );
+ assert_eq!(
+ full_join_partitioning(&left, &right, &a_equals_c()),
+ expected
+ );
+ }
+
+ #[test]
+ fn full_join_output_partitioning_coalesces_every_key_pair() {
+ let left = repartitioned(
+ int32_schema("a", "b"),
+ Partitioning::Hash(vec![col_at("a", 0), col_at("b", 1)], 4),
+ );
+ let right = repartitioned(
+ int32_schema("c", "d"),
+ Partitioning::Hash(vec![col_at("c", 0), col_at("d", 1)], 4),
+ );
+ let on = vec![
+ (col_at("a", 0), col_at("c", 0)),
+ (col_at("b", 1), col_at("d", 1)),
+ ];
+
+ let expected = Partitioning::Hash(
+ vec![
+ coalesced(col_at("a", 0), col_at("c", 2)),
+ coalesced(col_at("b", 1), col_at("d", 3)),
+ ],
+ 4,
+ );
+ assert_eq!(full_join_partitioning(&left, &right, &on), expected);
+ }
+
+ #[test]
+ fn
full_join_output_partitioning_is_unknown_for_different_range_split_points() {
+ let options = SortOptions::new(false, true);
+ let left = repartitioned(
+ int32_schema("a", "b"),
+ range_on(col_at("a", 0), options, &[10, 20, 30]),
+ );
+ let right = repartitioned(
+ int32_schema("c", "d"),
+ range_on(col_at("c", 0), options, &[10, 20, 40]),
+ );
+
+ assert_eq!(
+ full_join_partitioning(&left, &right, &a_equals_c()).to_string(),
+ "UnknownPartitioning(4)"
+ );
+ }
+
+ #[test]
+ fn full_join_output_partitioning_is_unknown_for_single_partition_inputs() {
+ let left = repartitioned(
+ int32_schema("a", "b"),
+ Partitioning::Hash(vec![col_at("a", 0)], 1),
+ );
+ let right = repartitioned(
+ int32_schema("c", "d"),
+ Partitioning::Hash(vec![col_at("c", 0)], 1),
+ );
+
+ assert_eq!(
+ full_join_partitioning(&left, &right, &a_equals_c()).to_string(),
+ "UnknownPartitioning(1)"
+ );
+ }
+
+ #[test]
+ fn
full_join_output_partitioning_is_unknown_when_keys_are_not_the_join_keys() {
+ let left = repartitioned(
+ int32_schema("a", "b"),
+ Partitioning::Hash(vec![col_at("b", 1)], 4),
+ );
+ let right = repartitioned(
+ int32_schema("c", "d"),
+ Partitioning::Hash(vec![col_at("c", 0)], 4),
+ );
+
+ assert_eq!(
+ full_join_partitioning(&left, &right, &a_equals_c()).to_string(),
+ "UnknownPartitioning(4)"
+ );
+ }
+
+ #[test]
+ fn full_join_output_partitioning_is_unknown_without_join_keys() {
+ let left = repartitioned(
+ int32_schema("a", "b"),
+ Partitioning::Hash(vec![col_at("a", 0)], 4),
+ );
+ let right = repartitioned(
+ int32_schema("c", "d"),
+ Partitioning::Hash(vec![col_at("c", 0)], 4),
+ );
+
+ assert_eq!(
+ full_join_partitioning(&left, &right, &[]).to_string(),
+ "UnknownPartitioning(4)"
+ );
+ }
+
+ #[test]
+ fn full_sort_merge_join_reports_coalesced_partitioning() {
Review Comment:
Added three tests for this through with_projection
--
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]
