diff --git a/datafusion/core/tests/repro_sort_api.rs b/datafusion/core/tests/repro_sort_api.rs new file mode 100644 index 0000000000000..336f21f15d403 --- /dev/null +++ b/datafusion/core/tests/repro_sort_api.rs @@ -0,0 +1,32 @@ +use datafusion::prelude::*; +use datafusion_common::Result; +use datafusion_expr::SortOptions; + +#[test] +fn test_sort_api_usage() -> Result<()> { + let expr = col("a"); + + // Old API: sort(asc, nulls_first) + // "True, False" -> Ascending, Nulls Last + let sort_expr = expr.clone().sort(true, false); + + assert_eq!(sort_expr.asc, true); + assert_eq!(sort_expr.nulls_first, false); + + // New API: sort_by with SortOptions + // Descending, Nulls First + let sort_expr = expr.clone().sort_by(SortOptions::new().desc().nulls_first()); + + assert_eq!(sort_expr.asc, false); + assert_eq!(sort_expr.nulls_first, true); + + // New API: Ascending, Nulls Last (default) + let sort_expr = expr.sort_by(SortOptions::new()); + + assert_eq!(sort_expr.asc, true); + assert_eq!(sort_expr.nulls_first, false); + + Ok(()) +} + + diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs index 09454795fd42d..4cf5087610231 100644 --- a/datafusion/expr/src/expr.rs +++ b/datafusion/expr/src/expr.rs @@ -1861,6 +1861,17 @@ impl Expr { Sort::new(self, asc, nulls_first) } + /// Create a sort configuration from an existing expression using `SortOptions`. + /// + /// ``` + /// # use datafusion_expr::{col, SortOptions}; + /// let sort_expr = col("foo").sort_by(SortOptions::new().desc().nulls_first()); + /// ``` + pub fn sort_by(self, options: crate::sort_options::SortOptions) -> Sort { + Sort::new(self, !options.descending, options.nulls_first) + } + + /// Return `IsTrue(Box(self))` pub fn is_true(self) -> Expr { Expr::IsTrue(Box::new(self)) diff --git a/datafusion/expr/src/lib.rs b/datafusion/expr/src/lib.rs index cb136229bf88d..3e2394eab467c 100644 --- a/datafusion/expr/src/lib.rs +++ b/datafusion/expr/src/lib.rs @@ -68,6 +68,7 @@ pub mod dml { pub mod planner; pub mod registry; pub mod simplify; +pub mod sort_options; pub mod sort_properties { pub use datafusion_expr_common::sort_properties::*; } @@ -128,6 +129,7 @@ pub use udaf::{ pub use udf::{ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl}; pub use udwf::{LimitEffect, ReversedUDWF, WindowUDF, WindowUDFImpl}; pub use window_frame::{WindowFrame, WindowFrameBound, WindowFrameUnits}; +pub use sort_options::SortOptions; #[cfg(test)] #[ctor::ctor] diff --git a/datafusion/expr/src/sort_options.rs b/datafusion/expr/src/sort_options.rs new file mode 100644 index 0000000000000..ed06c2c4ad0e4 --- /dev/null +++ b/datafusion/expr/src/sort_options.rs @@ -0,0 +1,53 @@ +use arrow::compute::SortOptions as ArrowSortOptions; + +/// Options for sorting. +/// +/// This struct implements a builder pattern for creating `SortOptions`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct SortOptions { + pub descending: bool, + pub nulls_first: bool, +} + +impl SortOptions { + /// Create a new `SortOptions` with default values (Ascending, Nulls Last). + pub fn new() -> Self { + Self { + descending: false, + nulls_first: false, + } + } + + /// Set sort order to descending. + pub fn desc(mut self) -> Self { + self.descending = true; + self + } + + /// Set sort order to ascending. + pub fn asc(mut self) -> Self { + self.descending = false; + self + } + + /// Set nulls to come first. + pub fn nulls_first(mut self) -> Self { + self.nulls_first = true; + self + } + + /// Set nulls to come last. + pub fn nulls_last(mut self) -> Self { + self.nulls_first = false; + self + } +} + +impl From for ArrowSortOptions { + fn from(options: SortOptions) -> Self { + ArrowSortOptions { + descending: options.descending, + nulls_first: options.nulls_first, + } + } +}