Merge branch 'apache:main' into fix-15428-map-distinct-tests

Author: PG1204

datafusion-examples/examples/udf/advanced_udaf.rs

@@ -23,8 +23,10 @@ use datafusion::{arrow::datatypes::DataType, logical_expr::Volatility};
 use std::sync::Arc;
 
 use arrow::array::{
-    ArrayRef, AsArray, Float32Array, PrimitiveArray, PrimitiveBuilder, UInt32Array,
+    Array, ArrayRef, AsArray, BooleanArray, Float32Array, PrimitiveArray,
+    PrimitiveBuilder, UInt32Array,
 };
+use arrow::buffer::NullBuffer;
 use arrow::datatypes::{ArrowNativeTypeOp, ArrowPrimitiveType, Float64Type, UInt32Type};
 use arrow::record_batch::RecordBatch;
 use arrow_schema::FieldRef;
@@ -237,7 +239,7 @@ impl GroupsAccumulator for GeometricMeanGroupsAccumulator {
         &mut self,
         values: &[ArrayRef],
         group_indices: &[usize],
-        opt_filter: Option<&arrow::array::BooleanArray>,
+        opt_filter: Option<&BooleanArray>,
         total_num_groups: usize,
     ) -> Result<()> {
         assert_eq!(values.len(), 1, "single argument to update_batch");
@@ -359,6 +361,43 @@ impl GroupsAccumulator for GeometricMeanGroupsAccumulator {
         ])
     }
 
+    fn convert_to_state(
+        &self,
+        values: &[ArrayRef],
+        opt_filter: Option<&BooleanArray>,
+    ) -> Result<Vec<ArrayRef>> {
+        assert_eq!(values.len(), 1, "single argument to convert_to_state");
+
+        let prods = values[0]
+            .as_primitive::<Float64Type>()
+            .clone()
+            .with_data_type(self.prod_data_type.clone());
+        let counts = UInt32Array::from_value(1, prods.len());
+
+        let filter_nulls = opt_filter.map(|filter| {
+            let validity = match filter.nulls() {
+                Some(nulls) => filter.values() & nulls.inner(),
+                None => filter.values().clone(),
+            };
+            NullBuffer::new(validity)
+        });
+        let nulls = NullBuffer::union(filter_nulls.as_ref(), prods.nulls());
+
+        let prods =
+            PrimitiveArray::<Float64Type>::new(prods.values().clone(), nulls.clone())
+                .with_data_type(self.prod_data_type.clone());
+        let counts = UInt32Array::new(counts.values().clone(), nulls);
+
+        Ok(vec![
+            Arc::new(prods) as ArrayRef,
+            Arc::new(counts) as ArrayRef,
+        ])
+    }
+
+    fn supports_convert_to_state(&self) -> bool {
+        true
+    }
+
     fn size(&self) -> usize {
         self.counts.capacity() * size_of::<u32>()
             + self.prods.capacity() * size_of::<Float64Type>()

datafusion/common-runtime/src/join_set.rs

@@ -21,10 +21,13 @@ use std::task::{Context, Poll};
 use tokio::runtime::Handle;
 use tokio::task::{AbortHandle, Id, JoinError, LocalSet};
 
-/// A wrapper around Tokio's JoinSet that forwards all API calls while optionally
+/// A wrapper around [Tokio's `JoinSet`] that forwards all API calls while optionally
 /// instrumenting spawned tasks and blocking closures with custom tracing behavior.
-/// If no tracer is injected via `trace_utils::set_tracer`, tasks and closures are executed
+/// If no tracer is injected via [`set_join_set_tracer`], tasks and closures are executed
 /// without any instrumentation.
+///
+/// [Tokio's `JoinSet`]: tokio::task::JoinSet
+/// [`set_join_set_tracer`]: crate::trace_utils::set_join_set_tracer
 #[derive(Debug)]
 pub struct JoinSet<T> {
     inner: tokio::task::JoinSet<T>,

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

@@ -265,8 +265,12 @@ impl ExecutionPlan for SinglePartitionMaintainsOrderExec {
         vec![&self.input]
     }
 
-    fn required_input_distribution(&self) -> Vec<Distribution> {
-        vec![Distribution::SinglePartition]
+    fn input_distribution_requirements(
+        &self,
+    ) -> datafusion_physical_plan::InputDistributionRequirements {
+        datafusion_physical_plan::InputDistributionRequirements::new(vec![
+            Distribution::SinglePartition,
+        ])
     }
 
     fn maintains_input_order(&self) -> Vec<bool> {
@@ -823,6 +827,49 @@ fn range_grouping_set_aggregate_rehashes_with_grouping_id() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn range_inner_hash_join_rehashes_incompatible_range_partitioning() -> Result<()> {
+    let left = parquet_exec_with_output_partitioning(range_partitioning(
+        "a",
+        [10, 20, 30],
+        SortOptions::default(),
+    )?);
+    let right = projection_exec_with_alias(
+        parquet_exec_with_output_partitioning(range_partitioning(
+            "a",
+            [10, 30, 40],
+            SortOptions::default(),
+        )?),
+        vec![
+            ("a".to_string(), "a1".to_string()),
+            ("b".to_string(), "b1".to_string()),
+        ],
+    );
+    let join_on = vec![(
+        Arc::new(Column::new_with_schema("a", &left.schema())?) as _,
+        Arc::new(Column::new_with_schema("a1", &right.schema())?) as _,
+    )];
+    let join = hash_join_exec(left, right, &join_on, &JoinType::Inner);
+
+    let plan = TestConfig::default()
+        .with_query_execution_partitions(4)
+        .to_plan(join, &DISTRIB_DISTRIB_SORT);
+
+    assert_plan!(
+        plan,
+        @r"
+    HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, a1@0)]
+      RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4
+        DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet
+      RepartitionExec: partitioning=Hash([a1@0], 4), input_partitions=4
+        ProjectionExec: expr=[a@0 as a1, b@1 as b1]
+          DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (30), (40)], 4), file_type=parquet
+    "
+    );
+
+    Ok(())
+}
+
 #[test]
 fn multi_hash_joins() -> Result<()> {
     let left = parquet_exec();

datafusion/core/tests/physical_optimizer/ensure_requirements.rs

@@ -974,8 +974,12 @@ impl ExecutionPlan for MockReqExec {
     fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
         vec![&self.input]
     }
-    fn required_input_distribution(&self) -> Vec<Distribution> {
-        vec![self.dist.clone()]
+    fn input_distribution_requirements(
+        &self,
+    ) -> datafusion_physical_plan::InputDistributionRequirements {
+        datafusion_physical_plan::InputDistributionRequirements::new(vec![
+            self.dist.clone(),
+        ])
     }
     fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
         vec![

datafusion/core/tests/physical_optimizer/projection_pushdown.rs

@@ -800,7 +800,9 @@ fn test_output_req_after_projection() -> Result<()> {
     if let Distribution::KeyPartitioned(vec) = after_optimize
         .downcast_ref::<OutputRequirementExec>()
         .unwrap()
-        .required_input_distribution()[0]
+        .input_distribution_requirements()
+        .child_distribution(0)
+        .unwrap()
         .clone()
     {
         assert!(

datafusion/core/tests/physical_optimizer/sanity_checker.rs

@@ -19,9 +19,9 @@ use insta::assert_snapshot;
 use std::sync::Arc;
 
 use crate::physical_optimizer::test_utils::{
-    bounded_window_exec, global_limit_exec, local_limit_exec, memory_exec,
-    projection_exec, repartition_exec, sort_exec, sort_expr, sort_expr_options,
-    sort_merge_join_exec, sort_preserving_merge_exec, union_exec,
+    bounded_window_exec, global_limit_exec, hash_join_exec, local_limit_exec,
+    memory_exec, projection_exec, repartition_exec, sort_exec, sort_expr,
+    sort_expr_options, sort_merge_join_exec, sort_preserving_merge_exec, union_exec,
 };
 
 use arrow::compute::SortOptions;
@@ -30,8 +30,8 @@ use datafusion::datasource::stream::{FileStreamProvider, StreamConfig, StreamTab
 use datafusion::prelude::{CsvReadOptions, SessionContext};
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::{JoinType, Result, ScalarValue};
-use datafusion_physical_expr::Partitioning;
 use datafusion_physical_expr::expressions::{Literal, col};
+use datafusion_physical_expr::{Partitioning, RangePartitioning, SplitPoint};
 use datafusion_physical_expr_common::sort_expr::LexOrdering;
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_optimizer::sanity_checker::SanityCheckPlan;
@@ -400,6 +400,49 @@ fn assert_sanity_check(plan: &Arc<dyn ExecutionPlan>, is_sane: bool) {
     );
 }
 
+fn range_partitioned_exec(
+    schema: &SchemaRef,
+    key: &str,
+    split_points: impl IntoIterator<Item = i32>,
+) -> Result<Arc<dyn ExecutionPlan>> {
+    let split_points = split_points
+        .into_iter()
+        .map(|value| SplitPoint::new(vec![ScalarValue::Int32(Some(value))]))
+        .collect();
+    let partitioning = Partitioning::Range(RangePartitioning::try_new(
+        [sort_expr(key, schema)].into(),
+        split_points,
+    )?);
+    RepartitionExec::try_new(memory_exec(schema), partitioning)
+        .map(|exec| Arc::new(exec) as Arc<dyn ExecutionPlan>)
+}
+
+#[test]
+fn test_partitioned_hash_join_requires_co_partitioned_children() -> Result<()> {
+    let schema = create_test_schema2();
+    let join_on = vec![(col("a", &schema)?, col("a", &schema)?)];
+
+    let compatible_join = hash_join_exec(
+        range_partitioned_exec(&schema, "a", [10])?,
+        range_partitioned_exec(&schema, "a", [10])?,
+        join_on.clone(),
+        None,
+        &JoinType::Inner,
+    )?;
+    assert_sanity_check(&compatible_join, true);
+
+    let incompatible_join = hash_join_exec(
+        range_partitioned_exec(&schema, "a", [10])?,
+        range_partitioned_exec(&schema, "a", [20])?,
+        join_on,
+        None,
+        &JoinType::Inner,
+    )?;
+    assert_sanity_check(&incompatible_join, false);
+
+    Ok(())
+}
+
 #[tokio::test]
 /// Tests that plan is valid when the sort requirements are satisfied.
 async fn test_bounded_window_agg_sort_requirement() -> Result<()> {

datafusion/core/tests/user_defined/user_defined_aggregates.rs

@@ -877,6 +877,22 @@ impl GroupsAccumulator for TestGroupsAccumulator {
         Ok(())
     }
 
+    fn convert_to_state(
+        &self,
+        values: &[ArrayRef],
+        _opt_filter: Option<&arrow::array::BooleanArray>,
+    ) -> Result<Vec<ArrayRef>> {
+        let len = values.first().map_or(0, |value| value.len());
+        Ok(vec![
+            Arc::new(PrimitiveArray::<UInt64Type>::from_value(self.result, len))
+                as ArrayRef,
+        ])
+    }
+
+    fn supports_convert_to_state(&self) -> bool {
+        true
+    }
+
     fn size(&self) -> usize {
         size_of::<u64>()
     }

datafusion/core/tests/user_defined/user_defined_plan.rs

@@ -708,8 +708,12 @@ impl ExecutionPlan for TopKExec {
         &self.cache
     }
 
-    fn required_input_distribution(&self) -> Vec<Distribution> {
-        vec![Distribution::SinglePartition]
+    fn input_distribution_requirements(
+        &self,
+    ) -> datafusion_physical_plan::InputDistributionRequirements {
+        datafusion_physical_plan::InputDistributionRequirements::new(vec![
+            Distribution::SinglePartition,
+        ])
     }
 
     fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {