Merge branch 'main' into physicalexpr-cse

# Conflicts:
#	datafusion/common/src/cse.rs
#	datafusion/optimizer/src/common_subexpr_eliminate.rs

Author: peter-toth

datafusion/core/src/physical_optimizer/join_selection.rs

@@ -140,20 +140,32 @@ fn swap_join_projection(
     left_schema_len: usize,
     right_schema_len: usize,
     projection: Option<&Vec<usize>>,
+    join_type: &JoinType,
 ) -> Option<Vec<usize>> {
-    projection.map(|p| {
-        p.iter()
-            .map(|i| {
-                // If the index is less than the left schema length, it is from the left schema, so we add the right schema length to it.
-                // Otherwise, it is from the right schema, so we subtract the left schema length from it.
-                if *i < left_schema_len {
-                    *i + right_schema_len
-                } else {
-                    *i - left_schema_len
-                }
-            })
-            .collect()
-    })
+    match join_type {
+        // For Anti/Semi join types, projection should remain unmodified,
+        // since these joins output schema remains the same after swap
+        JoinType::LeftAnti
+        | JoinType::LeftSemi
+        | JoinType::RightAnti
+        | JoinType::RightSemi => projection.cloned(),
+
+        _ => projection.map(|p| {
+            p.iter()
+                .map(|i| {
+                    // If the index is less than the left schema length, it is from
+                    // the left schema, so we add the right schema length to it.
+                    // Otherwise, it is from the right schema, so we subtract the left
+                    // schema length from it.
+                    if *i < left_schema_len {
+                        *i + right_schema_len
+                    } else {
+                        *i - left_schema_len
+                    }
+                })
+                .collect()
+        }),
+    }
 }
 
 /// This function swaps the inputs of the given join operator.
@@ -179,6 +191,7 @@ pub fn swap_hash_join(
             left.schema().fields().len(),
             right.schema().fields().len(),
             hash_join.projection.as_ref(),
+            hash_join.join_type(),
         ),
         partition_mode,
         hash_join.null_equals_null(),
@@ -1289,27 +1302,59 @@ mod tests_statistical {
         );
     }
 
+    #[rstest(
+        join_type, projection, small_on_right,
+        case::inner(JoinType::Inner, vec![1], true),
+        case::left(JoinType::Left, vec![1], true),
+        case::right(JoinType::Right, vec![1], true),
+        case::full(JoinType::Full, vec![1], true),
+        case::left_anti(JoinType::LeftAnti, vec![0], false),
+        case::left_semi(JoinType::LeftSemi, vec![0], false),
+        case::right_anti(JoinType::RightAnti, vec![0], true),
+        case::right_semi(JoinType::RightSemi, vec![0], true),
+    )]
     #[tokio::test]
-    async fn test_hash_join_swap_on_joins_with_projections() -> Result<()> {
+    async fn test_hash_join_swap_on_joins_with_projections(
+        join_type: JoinType,
+        projection: Vec<usize>,
+        small_on_right: bool,
+    ) -> Result<()> {
         let (big, small) = create_big_and_small();
+
+        let left = if small_on_right { &big } else { &small };
+        let right = if small_on_right { &small } else { &big };
+
+        let left_on = if small_on_right {
+            "big_col"
+        } else {
+            "small_col"
+        };
+        let right_on = if small_on_right {
+            "small_col"
+        } else {
+            "big_col"
+        };
+
         let join = Arc::new(HashJoinExec::try_new(
-            Arc::clone(&big),
-            Arc::clone(&small),
+            Arc::clone(left),
+            Arc::clone(right),
             vec![(
-                Arc::new(Column::new_with_schema("big_col", &big.schema())?),
-                Arc::new(Column::new_with_schema("small_col", &small.schema())?),
+                Arc::new(Column::new_with_schema(left_on, &left.schema())?),
+                Arc::new(Column::new_with_schema(right_on, &right.schema())?),
             )],
             None,
-            &JoinType::Inner,
-            Some(vec![1]),
+            &join_type,
+            Some(projection),
             PartitionMode::Partitioned,
             false,
         )?);
+
         let swapped = swap_hash_join(&join.clone(), PartitionMode::Partitioned)
             .expect("swap_hash_join must support joins with projections");
         let swapped_join = swapped.as_any().downcast_ref::<HashJoinExec>().expect(
             "ProjectionExec won't be added above if HashJoinExec contains embedded projection",
         );
+
         assert_eq!(swapped_join.projection, Some(vec![0_usize]));
         assert_eq!(swapped.schema().fields.len(), 1);
         assert_eq!(swapped.schema().fields[0].name(), "small_col");

datafusion/core/tests/fuzz_cases/equivalence/ordering.rs

@@ -16,8 +16,9 @@
 // under the License.
 
 use crate::fuzz_cases::equivalence::utils::{
-    create_random_schema, generate_table_for_eq_properties, is_table_same_after_sort,
-    TestScalarUDF,
+    convert_to_orderings, create_random_schema, create_test_schema_2,
+    generate_table_for_eq_properties, generate_table_for_orderings,
+    is_table_same_after_sort, TestScalarUDF,
 };
 use arrow_schema::SortOptions;
 use datafusion_common::{DFSchema, Result};
@@ -158,3 +159,66 @@ fn test_ordering_satisfy_with_equivalence_complex_random() -> Result<()> {
 
     Ok(())
 }
+
+// This test checks given a table is ordered with `[a ASC, b ASC, c ASC, d ASC]` and `[a ASC, c ASC, b ASC, d ASC]`
+// whether the table is also ordered with `[a ASC, b ASC, d ASC]` and `[a ASC, c ASC, d ASC]`
+// Since these orderings cannot be deduced, these orderings shouldn't be satisfied by the table generated.
+// For background see discussion: https://github.com/apache/datafusion/issues/12700#issuecomment-2411134296
+#[test]
+fn test_ordering_satisfy_on_data() -> Result<()> {
+    let schema = create_test_schema_2()?;
+    let col_a = &col("a", &schema)?;
+    let col_b = &col("b", &schema)?;
+    let col_c = &col("c", &schema)?;
+    let col_d = &col("d", &schema)?;
+
+    let option_asc = SortOptions {
+        descending: false,
+        nulls_first: false,
+    };
+
+    let orderings = vec![
+        // [a ASC, b ASC, c ASC, d ASC]
+        vec![
+            (col_a, option_asc),
+            (col_b, option_asc),
+            (col_c, option_asc),
+            (col_d, option_asc),
+        ],
+        // [a ASC, c ASC, b ASC, d ASC]
+        vec![
+            (col_a, option_asc),
+            (col_c, option_asc),
+            (col_b, option_asc),
+            (col_d, option_asc),
+        ],
+    ];
+    let orderings = convert_to_orderings(&orderings);
+
+    let batch = generate_table_for_orderings(orderings, schema, 1000, 10)?;
+
+    // [a ASC, c ASC, d ASC] cannot be deduced
+    let ordering = vec![
+        (col_a, option_asc),
+        (col_c, option_asc),
+        (col_d, option_asc),
+    ];
+    let ordering = convert_to_orderings(&[ordering])[0].clone();
+    assert!(!is_table_same_after_sort(ordering, batch.clone())?);
+
+    // [a ASC, b ASC, d ASC] cannot be deduced
+    let ordering = vec![
+        (col_a, option_asc),
+        (col_b, option_asc),
+        (col_d, option_asc),
+    ];
+    let ordering = convert_to_orderings(&[ordering])[0].clone();
+    assert!(!is_table_same_after_sort(ordering, batch.clone())?);
+
+    // [a ASC, b ASC] can be deduced
+    let ordering = vec![(col_a, option_asc), (col_b, option_asc)];
+    let ordering = convert_to_orderings(&[ordering])[0].clone();
+    assert!(is_table_same_after_sort(ordering, batch.clone())?);
+
+    Ok(())
+}

datafusion/core/tests/fuzz_cases/equivalence/utils.rs

@@ -15,23 +15,29 @@
 // specific language governing permissions and limitations
 // under the License.
 //
-// use datafusion_physical_expr::expressions::{col, Column};
 use datafusion::physical_plan::expressions::col;
 use datafusion::physical_plan::expressions::Column;
 use datafusion_physical_expr::{ConstExpr, EquivalenceProperties, PhysicalSortExpr};
 use std::any::Any;
+use std::cmp::Ordering;
 use std::sync::Arc;
 
 use arrow::compute::{lexsort_to_indices, SortColumn};
 use arrow::datatypes::{DataType, Field, Schema};
-use arrow_array::{ArrayRef, Float32Array, Float64Array, RecordBatch, UInt32Array};
+use arrow_array::{
+    ArrayRef, Float32Array, Float64Array, PrimitiveArray, RecordBatch, UInt32Array,
+};
 use arrow_schema::{SchemaRef, SortOptions};
+use datafusion_common::utils::{
+    compare_rows, get_record_batch_at_indices, get_row_at_idx,
+};
 use datafusion_common::{exec_err, plan_datafusion_err, DataFusionError, Result};
-
 use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
 use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
 use datafusion_physical_expr::equivalence::{EquivalenceClass, ProjectionMapping};
 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
+use datafusion_physical_expr_common::sort_expr::{LexOrdering, LexOrderingRef};
+
 use itertools::izip;
 use rand::prelude::*;
 
@@ -67,7 +73,7 @@ pub fn output_schema(
 }
 
 // Generate a schema which consists of 6 columns (a, b, c, d, e, f)
-fn create_test_schema_2() -> Result<SchemaRef> {
+pub fn create_test_schema_2() -> Result<SchemaRef> {
     let a = Field::new("a", DataType::Float64, true);
     let b = Field::new("b", DataType::Float64, true);
     let c = Field::new("c", DataType::Float64, true);
@@ -374,6 +380,114 @@ pub fn generate_table_for_eq_properties(
     Ok(RecordBatch::try_from_iter(res)?)
 }
 
+// Generate a table that satisfies the given orderings;
+pub fn generate_table_for_orderings(
+    mut orderings: Vec<LexOrdering>,
+    schema: SchemaRef,
+    n_elem: usize,
+    n_distinct: usize,
+) -> Result<RecordBatch> {
+    let mut rng = StdRng::seed_from_u64(23);
+
+    assert!(!orderings.is_empty());
+    // Sort the inner vectors by their lengths (longest first)
+    orderings.sort_by_key(|v| std::cmp::Reverse(v.len()));
+
+    let arrays = schema
+        .fields
+        .iter()
+        .map(|field| {
+            (
+                field.name(),
+                generate_random_f64_array(n_elem, n_distinct, &mut rng),
+            )
+        })
+        .collect::<Vec<_>>();
+    let batch = RecordBatch::try_from_iter(arrays)?;
+
+    // Sort batch according to first ordering expression
+    let sort_columns = get_sort_columns(&batch, &orderings[0])?;
+    let sort_indices = lexsort_to_indices(&sort_columns, None)?;
+    let mut batch = get_record_batch_at_indices(&batch, &sort_indices)?;
+
+    // prune out rows that is invalid according to remaining orderings.
+    for ordering in orderings.iter().skip(1) {
+        let sort_columns = get_sort_columns(&batch, ordering)?;
+
+        // Collect sort options and values into separate vectors.
+        let (sort_options, sort_col_values): (Vec<_>, Vec<_>) = sort_columns
+            .into_iter()
+            .map(|sort_col| (sort_col.options.unwrap(), sort_col.values))
+            .unzip();
+
+        let mut cur_idx = 0;
+        let mut keep_indices = vec![cur_idx as u32];
+        for next_idx in 1..batch.num_rows() {
+            let cur_row = get_row_at_idx(&sort_col_values, cur_idx)?;
+            let next_row = get_row_at_idx(&sort_col_values, next_idx)?;
+
+            if compare_rows(&cur_row, &next_row, &sort_options)? != Ordering::Greater {
+                // next row satisfies ordering relation given, compared to the current row.
+                keep_indices.push(next_idx as u32);
+                cur_idx = next_idx;
+            }
+        }
+        // Only keep valid rows, that satisfies given ordering relation.
+        batch = get_record_batch_at_indices(
+            &batch,
+            &PrimitiveArray::from_iter_values(keep_indices),
+        )?;
+    }
+
+    Ok(batch)
+}
+
+// Convert each tuple to PhysicalSortExpr
+pub fn convert_to_sort_exprs(
+    in_data: &[(&Arc<dyn PhysicalExpr>, SortOptions)],
+) -> Vec<PhysicalSortExpr> {
+    in_data
+        .iter()
+        .map(|(expr, options)| PhysicalSortExpr {
+            expr: Arc::clone(*expr),
+            options: *options,
+        })
+        .collect()
+}
+
+// Convert each inner tuple to PhysicalSortExpr
+pub fn convert_to_orderings(
+    orderings: &[Vec<(&Arc<dyn PhysicalExpr>, SortOptions)>],
+) -> Vec<Vec<PhysicalSortExpr>> {
+    orderings
+        .iter()
+        .map(|sort_exprs| convert_to_sort_exprs(sort_exprs))
+        .collect()
+}
+
+// Utility function to generate random f64 array
+fn generate_random_f64_array(
+    n_elems: usize,
+    n_distinct: usize,
+    rng: &mut StdRng,
+) -> ArrayRef {
+    let values: Vec<f64> = (0..n_elems)
+        .map(|_| rng.gen_range(0..n_distinct) as f64 / 2.0)
+        .collect();
+    Arc::new(Float64Array::from_iter_values(values))
+}
+
+// Helper function to get sort columns from a batch
+fn get_sort_columns(
+    batch: &RecordBatch,
+    ordering: LexOrderingRef,
+) -> Result<Vec<SortColumn>> {
+    ordering
+        .iter()
+        .map(|expr| expr.evaluate_to_sort_column(batch))
+        .collect::<Result<Vec<_>>>()
+}
+
 #[derive(Debug, Clone)]
 pub struct TestScalarUDF {
     pub(crate) signature: Signature,

datafusion/expr/src/udf_docs.rs

@@ -147,23 +147,29 @@ impl DocumentationBuilder {
 
     /// Add a standard "expression" argument to the documentation
     ///
-    /// This is similar to  [`Self::with_argument`] except that  a standard
-    /// description is appended to the end: `"Can be a constant, column, or
-    /// function, and any combination of arithmetic operators."`
-    ///
-    /// The argument is rendered like
+    /// The argument is rendered like below if Some() is passed through:
     ///
     /// ```text
     /// <arg_name>:
     ///   <expression_type> expression to operate on. Can be a constant, column, or function, and any combination of operators.
     /// ```
+    ///
+    /// The argument is rendered like below if None is passed through:
+    ///
+    ///  ```text
+    /// <arg_name>:
+    ///   The expression to operate on. Can be a constant, column, or function, and any combination of operators.
+    /// ```
     pub fn with_standard_argument(
         self,
         arg_name: impl Into<String>,
-        expression_type: impl AsRef<str>,
+        expression_type: Option<&str>,
     ) -> Self {
-        let expression_type = expression_type.as_ref();
-        self.with_argument(arg_name, format!("{expression_type} expression to operate on. Can be a constant, column, or function, and any combination of operators."))
+        let description = format!(
+            "{} expression to operate on. Can be a constant, column, or function, and any combination of operators.",
+            expression_type.unwrap_or("The")
+        );
+        self.with_argument(arg_name, description)
     }
 
     pub fn with_related_udf(mut self, related_udf: impl Into<String>) -> Self {