Feature Add Monotonic Definition

datafusion/core/src/physical_optimizer/enforce_sorting.rs

@@ -665,7 +665,9 @@ mod tests {
     use crate::prelude::{SessionConfig, SessionContext};
     use crate::test::{csv_exec_ordered, csv_exec_sorted, stream_exec_ordered};
     use datafusion_physical_optimizer::test_utils::{
-        aggregate_exec, bounded_window_exec, check_integrity, coalesce_batches_exec,
+        aggregate_exec, aggregate_exec_monotonic, aggregate_exec_non_monotonic,
+        bounded_window_exec, bounded_window_exec_non_monotonic,
+        bounded_window_exec_with_partition, check_integrity, coalesce_batches_exec,
         coalesce_partitions_exec, filter_exec, global_limit_exec, hash_join_exec,
         limit_exec, local_limit_exec, memory_exec, repartition_exec, sort_exec,
         sort_expr, sort_expr_options, sort_merge_join_exec, sort_preserving_merge_exec,
@@ -707,6 +709,17 @@ mod tests {
         Ok(schema)
     }
 
+    // Generate a schema which consists of 5 columns (a, b, c, d, e) of Uint64
+    fn create_test_schema4() -> Result<SchemaRef> {
+        let a = Field::new("a", DataType::UInt64, true);
+        let b = Field::new("b", DataType::UInt64, false);
+        let c = Field::new("c", DataType::UInt64, true);
+        let d = Field::new("d", DataType::UInt64, false);
+        let e = Field::new("e", DataType::UInt64, false);
+        let schema = Arc::new(Schema::new(vec![a, b, c, d, e]));
+        Ok(schema)
+    }
+
     /// Runs the sort enforcement optimizer and asserts the plan
     /// against the original and expected plans
     ///
@@ -1010,6 +1023,63 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_aggregate_monotonic() -> Result<()> {
+        let schema = create_test_schema4()?;
+        let source = memory_exec(&schema);
+        let sort_exprs = vec![sort_expr("a", &schema)];
+        let sort = sort_exec(sort_exprs.clone(), source);
+
+        let aggregate = aggregate_exec_monotonic(sort);
+        let sort_exprs = LexOrdering::new(vec![sort_expr("count", &aggregate.schema())]);
+        let physical_plan: Arc<dyn ExecutionPlan> =
+            Arc::new(SortExec::new(sort_exprs.clone(), aggregate)) as _;
+
+        let expected_input = [
+            "SortExec: expr=[count@0 ASC], preserve_partitioning=[false]",
+            "  AggregateExec: mode=Final, gby=[], aggr=[count]",
+            "    SortExec: expr=[a@0 ASC], preserve_partitioning=[false]",
+            "      MemoryExec: partitions=1, partition_sizes=[0]",
+        ];
+
+        let expected_optimized = [
+            "AggregateExec: mode=Final, gby=[], aggr=[count]",
+            "  MemoryExec: partitions=1, partition_sizes=[0]",
+        ];
+        assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_aggregate_non_monotonic() -> Result<()> {
+        let schema = create_test_schema4()?;
+        let source = memory_exec(&schema);
+        let sort_exprs = vec![sort_expr("a", &schema)];
+        let sort = sort_exec(sort_exprs.clone(), source);
+
+        let aggregate = aggregate_exec_non_monotonic(sort);
+        let sort_exprs = LexOrdering::new(vec![sort_expr("avg", &aggregate.schema())]);
+        let physical_plan: Arc<dyn ExecutionPlan> =
+            Arc::new(SortExec::new(sort_exprs.clone(), aggregate)) as _;
+
+        let expected_input = [
+            "SortExec: expr=[avg@0 ASC], preserve_partitioning=[false]",
+            "  AggregateExec: mode=Final, gby=[], aggr=[avg]",
+            "    SortExec: expr=[a@0 ASC], preserve_partitioning=[false]",
+            "      MemoryExec: partitions=1, partition_sizes=[0]",
+        ];
+
+        let expected_optimized = [
+            "SortExec: expr=[avg@0 ASC], preserve_partitioning=[false]",
+            "  AggregateExec: mode=Final, gby=[], aggr=[avg]",
+            "    MemoryExec: partitions=1, partition_sizes=[0]",
+        ];
+        assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_remove_unnecessary_sort4() -> Result<()> {
         let schema = create_test_schema()?;
@@ -1758,6 +1828,207 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_bounded_window_monotonic_sort() -> Result<()> {
+        let schema = create_test_schema()?;
+        let sort_exprs = vec![sort_expr_options(
+            "nullable_col",
+            &schema,
+            SortOptions {
+                descending: true,
+                nulls_first: false,
+            },
+        )];
+        let source = parquet_exec_sorted(&schema, sort_exprs.clone());
+        let sort = sort_exec(sort_exprs.clone(), source);
+
+        let bounded_window =
+            bounded_window_exec("nullable_col", sort_exprs.clone(), sort);
+        let output_schema = bounded_window.schema();
+        let sort_exprs2 = vec![sort_expr_options(
+            "count",
+            &output_schema,
+            SortOptions {
+                descending: false,
+                nulls_first: false,
+            },
+        )];
+        let physical_plan = sort_exec(sort_exprs2.clone(), bounded_window);
+
+        let expected_input = [
+            "SortExec: expr=[count@2 ASC NULLS LAST], preserve_partitioning=[false]",
+            "  BoundedWindowAggExec: wdw=[count: Ok(Field { name: \"count\", data_type: Int64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Range, start_bound: Preceding(NULL), end_bound: CurrentRow, is_causal: false }], mode=[Sorted]",
+            "    SortExec: expr=[nullable_col@0 DESC NULLS LAST], preserve_partitioning=[false]",
+            "      ParquetExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 DESC NULLS LAST]",
+        ];
+        let expected_optimized = [
+            "BoundedWindowAggExec: wdw=[count: Ok(Field { name: \"count\", data_type: Int64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Range, start_bound: Preceding(NULL), end_bound: CurrentRow, is_causal: false }], mode=[Sorted]",
+            "  ParquetExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 DESC NULLS LAST]",
+        ];
+        assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_bounded_plain_window_monotonic_sort_with_partitions() -> Result<()> {
+        let schema = create_test_schema()?;
+        let sort_exprs = vec![sort_expr_options(
+            "nullable_col",
+            &schema,
+            SortOptions {
+                descending: true,
+                nulls_first: false,
+            },
+        )];
+        let source = parquet_exec_sorted(&schema, sort_exprs.clone());
+        let sort = sort_exec(sort_exprs.clone(), source);
+        let partition_bys = &[col("nullable_col", &schema)?];
+
+        let bounded_window = bounded_window_exec_with_partition(
+            "nullable_col",
+            sort_exprs.clone(),
+            partition_bys,
+            sort,
+            false,
+        );
+        let output_schema = bounded_window.schema();
+        let sort_exprs2 = vec![
+            sort_expr_options(
+                "nullable_col",
+                &output_schema,
+                SortOptions {
+                    descending: true,
+                    nulls_first: false,
+                },
+            ),
+            sort_expr_options(
+                "count",
+                &output_schema,
+                SortOptions {
+                    descending: false,
+                    nulls_first: false,
+                },
+            ),
+        ];
+        let physical_plan = sort_exec(sort_exprs2.clone(), bounded_window);
+
+        let expected_input = [
+            "SortExec: expr=[nullable_col@0 DESC NULLS LAST, count@2 ASC NULLS LAST], preserve_partitioning=[false]",
+            "  BoundedWindowAggExec: wdw=[count: Ok(Field { name: \"count\", data_type: Int64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Range, start_bound: Preceding(NULL), end_bound: CurrentRow, is_causal: false }], mode=[Sorted]",
+            "    SortExec: expr=[nullable_col@0 DESC NULLS LAST], preserve_partitioning=[false]",
+            "      ParquetExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 DESC NULLS LAST]",
+        ];
+        let expected_optimized = [
+            "BoundedWindowAggExec: wdw=[count: Ok(Field { name: \"count\", data_type: Int64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Range, start_bound: Preceding(NULL), end_bound: CurrentRow, is_causal: false }], mode=[Sorted]",
+            "  ParquetExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 DESC NULLS LAST]",
+        ];
+        assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_bounded_sliding_window_monotonic_sort_with_partitions() -> Result<()> {
+        let schema = create_test_schema()?;
+        let sort_exprs = vec![sort_expr_options(
+            "nullable_col",
+            &schema,
+            SortOptions {
+                descending: true,
+                nulls_first: false,
+            },
+        )];
+        let source = parquet_exec_sorted(&schema, sort_exprs.clone());
+        let sort = sort_exec(sort_exprs.clone(), source);
+        let partition_bys = &[col("nullable_col", &schema)?];
+
+        let bounded_window = bounded_window_exec_with_partition(
+            "nullable_col",
+            sort_exprs.clone(),
+            partition_bys,
+            sort,
+            true,
+        );
+        let output_schema = bounded_window.schema();
+        let sort_exprs2 = vec![
+            sort_expr_options(
+                "nullable_col",
+                &output_schema,
+                SortOptions {
+                    descending: true,
+                    nulls_first: false,
+                },
+            ),
+            sort_expr_options(
+                "count",
+                &output_schema,
+                SortOptions {
+                    descending: false,
+                    nulls_first: false,
+                },
+            ),
+        ];
+        let physical_plan = sort_exec(sort_exprs2.clone(), bounded_window);
+
+        let expected_input = [
+            "SortExec: expr=[nullable_col@0 DESC NULLS LAST, count@2 ASC NULLS LAST], preserve_partitioning=[false]",
+            "  BoundedWindowAggExec: wdw=[count: Ok(Field { name: \"count\", data_type: Int64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Range, start_bound: CurrentRow, end_bound: Following(NULL), is_causal: false }], mode=[Sorted]",
+            "    SortExec: expr=[nullable_col@0 DESC NULLS LAST], preserve_partitioning=[false]",
+            "      ParquetExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 DESC NULLS LAST]",
+        ];
+        let expected_optimized = [
+            "BoundedWindowAggExec: wdw=[count: Ok(Field { name: \"count\", data_type: Int64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Range, start_bound: CurrentRow, end_bound: Following(NULL), is_causal: false }], mode=[Sorted]",
+            "  ParquetExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 DESC NULLS LAST]",
+        ];
+        assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_bounded_window_non_monotonic_sort() -> Result<()> {
+        let schema = create_test_schema4()?;
+        let sort_exprs = vec![sort_expr_options(
+            "a",
+            &schema,
+            SortOptions {
+                descending: true,
+                nulls_first: false,
+            },
+        )];
+        let source = parquet_exec_sorted(&schema, sort_exprs.clone());
+        let sort = sort_exec(sort_exprs.clone(), source);
+
+        let bounded_window =
+            bounded_window_exec_non_monotonic("a", sort_exprs.clone(), sort);
+        let output_schema = bounded_window.schema();
+        let sort_exprs2 = vec![sort_expr_options(
+            "avg",
+            &output_schema,
+            SortOptions {
+                descending: false,
+                nulls_first: false,
+            },
+        )];
+        let physical_plan = sort_exec(sort_exprs2.clone(), bounded_window);
+
+        let expected_input = [
+            "SortExec: expr=[avg@5 ASC NULLS LAST], preserve_partitioning=[false]",
+            "  BoundedWindowAggExec: wdw=[avg: Ok(Field { name: \"avg\", data_type: Float64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Range, start_bound: Preceding(NULL), end_bound: CurrentRow, is_causal: false }], mode=[Sorted]",
+            "    SortExec: expr=[a@0 DESC NULLS LAST], preserve_partitioning=[false]",
+            "      ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 DESC NULLS LAST]",
+        ];
+        let expected_optimized = [
+            "SortExec: expr=[avg@5 ASC NULLS LAST], preserve_partitioning=[false]",
+            "  BoundedWindowAggExec: wdw=[avg: Ok(Field { name: \"avg\", data_type: Float64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Range, start_bound: Preceding(NULL), end_bound: CurrentRow, is_causal: false }], mode=[Sorted]",
+            "    ParquetExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 DESC NULLS LAST]",
+        ];
+        assert_optimized!(expected_input, expected_optimized, physical_plan, true);
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_window_multi_path_sort2() -> Result<()> {
         let schema = create_test_schema()?;

datafusion/expr/src/test/function_stub.rs

@@ -278,6 +278,10 @@ impl AggregateUDFImpl for Count {
     fn reverse_expr(&self) -> ReversedUDAF {
         ReversedUDAF::Identical
     }
+
+    fn is_monotonic(&self) -> Option<bool> {
+        Some(true)
+    }
 }
 
 create_func!(Min, min_udaf);
@@ -363,6 +367,9 @@ impl AggregateUDFImpl for Min {
     fn is_descending(&self) -> Option<bool> {
         Some(false)
     }
+    fn is_monotonic(&self) -> Option<bool> {
+        Some(false)
+    }
 }
 
 create_func!(Max, max_udaf);
@@ -448,6 +455,9 @@ impl AggregateUDFImpl for Max {
     fn is_descending(&self) -> Option<bool> {
         Some(true)
     }
+    fn is_monotonic(&self) -> Option<bool> {
+        Some(true)
+    }
 }
 
 /// Testing stub implementation of avg aggregate

datafusion/expr/src/udaf.rs

@@ -389,7 +389,7 @@ pub trait AggregateUDFImpl: Debug + Send + Sync {
 
     /// Whether the aggregate function is nullable.
     ///
-    /// Nullable means that that the function could return `null` for any inputs.
+    /// Nullable means that the function could return `null` for any inputs.
     /// For example, aggregate functions like `COUNT` always return a non null value
     /// but others like `MIN` will return `NULL` if there is nullable input.
     /// Note that if the function is declared as *not* nullable, make sure the [`AggregateUDFImpl::default_value`] is `non-null`
@@ -635,6 +635,18 @@ pub trait AggregateUDFImpl: Debug + Send + Sync {
     fn documentation(&self) -> Option<&Documentation> {
         None
     }
+
+    /// Indicates whether the aggregation function is monotonic as a set function. A set
+    /// function is monotonically increasing if its value increases as its argument grows
+    /// (as a set). Formally, `f` is a monotonically increasing set function if `f(S) >= f(T)`
+    /// whenever `S` is a superset of `T`.
+    ///
+    /// Returns None if the function is not monotonic.
+    /// If the function is monotonically decreasing returns Some(false) e.g. Min
+    /// If the function is monotonically increasing returns Some(true) e.g. Max
+    fn is_monotonic(&self) -> Option<bool> {
+        None
+    }
 }
 
 impl PartialEq for dyn AggregateUDFImpl {

datafusion/expr/src/window_frame.rs

@@ -291,6 +291,13 @@ impl WindowFrame {
             && (self.end_bound.is_unbounded()
                 || self.end_bound == WindowFrameBound::CurrentRow)
     }
+
+    /// Is the window frame ever-expanding (it always grows in the superset sense).
+    /// Useful when understanding if set-monotonicity properties of functions can
+    /// be exploited.
+    pub fn is_ever_expanding(&self) -> bool {
+        self.start_bound.is_unbounded()
+    }
 }
 
 /// There are five ways to describe starting and ending frame boundaries:

datafusion/functions-aggregate/src/count.rs

@@ -351,6 +351,10 @@ impl AggregateUDFImpl for Count {
     fn documentation(&self) -> Option<&Documentation> {
         self.doc()
     }
+
+    fn is_monotonic(&self) -> Option<bool> {
+        Some(true)
+    }
 }
 
 #[derive(Debug)]