diff --git a/arrow-row/src/lib.rs b/arrow-row/src/lib.rs index ac859e90e126..e496419ff842 100644 --- a/arrow-row/src/lib.rs +++ b/arrow-row/src/lib.rs @@ -1492,10 +1492,7 @@ mod tests { let converter = RowConverter::new(vec![SortField::new_with_options( DataType::Boolean, - SortOptions { - descending: true, - nulls_first: false, - }, + SortOptions::default().desc().with_nulls_first(false), )]) .unwrap(); @@ -1613,10 +1610,7 @@ mod tests { let converter = RowConverter::new(vec![SortField::new_with_options( DataType::Binary, - SortOptions { - descending: true, - nulls_first: false, - }, + SortOptions::default().desc().with_nulls_first(false), )]) .unwrap(); let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap(); @@ -1695,10 +1689,7 @@ mod tests { let converter = RowConverter::new(vec![SortField::new_with_options( a.data_type().clone(), - SortOptions { - descending: true, - nulls_first: false, - }, + SortOptions::default().desc().with_nulls_first(false), )]) .unwrap(); @@ -1713,10 +1704,7 @@ mod tests { let converter = RowConverter::new(vec![SortField::new_with_options( a.data_type().clone(), - SortOptions { - descending: true, - nulls_first: true, - }, + SortOptions::default().desc().with_nulls_first(true), )]) .unwrap(); @@ -1889,10 +1877,7 @@ mod tests { back[0].to_data().validate_full().unwrap(); assert_eq!(&back[0], &list); - let options = SortOptions { - descending: false, - nulls_first: false, - }; + let options = SortOptions::default().asc().with_nulls_first(false); let field = SortField::new_with_options(d.clone(), options); let converter = RowConverter::new(vec![field]).unwrap(); let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap(); @@ -1909,10 +1894,7 @@ mod tests { back[0].to_data().validate_full().unwrap(); assert_eq!(&back[0], &list); - let options = SortOptions { - descending: true, - nulls_first: false, - }; + let options = SortOptions::default().desc().with_nulls_first(false); let field = SortField::new_with_options(d.clone(), options); let converter = RowConverter::new(vec![field]).unwrap(); let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap(); @@ -1929,10 +1911,7 @@ mod tests { back[0].to_data().validate_full().unwrap(); assert_eq!(&back[0], &list); - let options = SortOptions { - descending: true, - nulls_first: true, - }; + let options = SortOptions::default().desc().with_nulls_first(true); let field = SortField::new_with_options(d, options); let converter = RowConverter::new(vec![field]).unwrap(); let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap(); @@ -1992,10 +1971,7 @@ mod tests { // null // [[1, 2]] // ] - let options = SortOptions { - descending: false, - nulls_first: true, - }; + let options = SortOptions::default().asc().with_nulls_first(true); let field = SortField::new_with_options(d.clone(), options); let converter = RowConverter::new(vec![field]).unwrap(); let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap(); @@ -2011,10 +1987,7 @@ mod tests { back[0].to_data().validate_full().unwrap(); assert_eq!(&back[0], &list); - let options = SortOptions { - descending: true, - nulls_first: true, - }; + let options = SortOptions::default().desc().with_nulls_first(true); let field = SortField::new_with_options(d.clone(), options); let converter = RowConverter::new(vec![field]).unwrap(); let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap(); @@ -2030,10 +2003,7 @@ mod tests { back[0].to_data().validate_full().unwrap(); assert_eq!(&back[0], &list); - let options = SortOptions { - descending: true, - nulls_first: false, - }; + let options = SortOptions::default().desc().with_nulls_first(false); let field = SortField::new_with_options(d, options); let converter = RowConverter::new(vec![field]).unwrap(); let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap(); diff --git a/arrow-schema/src/lib.rs b/arrow-schema/src/lib.rs index d513ca7dd526..d06382fbcdf7 100644 --- a/arrow-schema/src/lib.rs +++ b/arrow-schema/src/lib.rs @@ -19,7 +19,9 @@ //! Arrow logical types mod datatype; + pub use datatype::*; +use std::fmt::Display; mod datatype_parse; mod error; pub use error::*; @@ -35,6 +37,42 @@ use std::ops; pub mod ffi; /// Options that define the sort order of a given column +/// +/// The default sorts equivalently to of `ASC NULLS FIRST` in SQL (i.e. +/// ascending order with nulls sorting before any other values). +/// +/// # Example creation +/// ``` +/// # use arrow_schema::SortOptions; +/// // configure using explicit initialization +/// let options = SortOptions { +/// descending: false, +/// nulls_first: true, +/// }; +/// // Default is ASC NULLs First +/// assert_eq!(options, SortOptions::default()); +/// assert_eq!(options.to_string(), "ASC NULLS FIRST"); +/// +/// // Configure using builder APIs +/// let options = SortOptions::default() +/// .desc() +/// .nulls_first(); +/// assert_eq!(options.to_string(), "DESC NULLS FIRST"); +/// +/// // configure using explicit field values +/// let options = SortOptions::default() +/// .with_descending(false) +/// .with_nulls_first(false); +/// assert_eq!(options.to_string(), "ASC NULLS LAST"); +/// ``` +/// +/// # Example operations +/// It is also possible to negate the sort options using the `!` operator. +/// ``` +/// use arrow_schema::SortOptions; +/// let options = !SortOptions::default(); +/// assert_eq!(options.to_string(), "DESC NULLS LAST"); +/// ``` #[derive(Clone, Hash, Copy, Debug, Eq, PartialEq, Ord, PartialOrd)] pub struct SortOptions { /// Whether to sort in descending order @@ -43,6 +81,76 @@ pub struct SortOptions { pub nulls_first: bool, } +impl Display for SortOptions { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + if self.descending { + write!(f, "DESC")?; + } else { + write!(f, "ASC")?; + } + if self.nulls_first { + write!(f, " NULLS FIRST")?; + } else { + write!(f, " NULLS LAST")?; + } + Ok(()) + } +} + +impl SortOptions { + /// Create a new `SortOptions` struct + pub fn new(descending: bool, nulls_first: bool) -> Self { + Self { + descending, + nulls_first, + } + } + + /// Set this sort options to sort in descending order + /// + /// See [Self::with_descending] to explicitly set the underlying field + pub fn desc(mut self) -> Self { + self.descending = true; + self + } + + /// Set this sort options to sort in ascending order + /// + /// See [Self::with_descending] to explicitly set the underlying field + pub fn asc(mut self) -> Self { + self.descending = false; + self + } + + /// Set this sort options to sort nulls first + /// + /// See [Self::with_nulls_first] to explicitly set the underlying field + pub fn nulls_first(mut self) -> Self { + self.nulls_first = true; + self + } + + /// Set this sort options to sort nulls last + /// + /// See [Self::with_nulls_first] to explicitly set the underlying field + pub fn nulls_last(mut self) -> Self { + self.nulls_first = false; + self + } + + /// Set this sort options to sort descending if argument is true + pub fn with_descending(mut self, descending: bool) -> Self { + self.descending = descending; + self + } + + /// Set this sort options to sort nulls first if argument is true + pub fn with_nulls_first(mut self, nulls_first: bool) -> Self { + self.nulls_first = nulls_first; + self + } +} + impl Default for SortOptions { fn default() -> Self { Self {