Skip to content

Commit

Permalink
Add additional documentation and builder APIs to SortOptions (#6441)
Browse files Browse the repository at this point in the history
* Minor: Add additional documentation and builder APIs to `SortOptions`

* Port some uses

* Update defaults

* Add nulls_first() and nulls_last() and more examples
  • Loading branch information
alamb committed Sep 25, 2024
1 parent 62825b2 commit 6137e91
Show file tree
Hide file tree
Showing 2 changed files with 118 additions and 40 deletions.
50 changes: 10 additions & 40 deletions arrow-row/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();

Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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();

Expand All @@ -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();

Expand Down Expand Up @@ -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();
Expand All @@ -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();
Expand All @@ -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();
Expand Down Expand Up @@ -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();
Expand All @@ -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();
Expand All @@ -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();
Expand Down
108 changes: 108 additions & 0 deletions arrow-schema/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,9 @@
//! Arrow logical types

mod datatype;

pub use datatype::*;
use std::fmt::Display;
mod datatype_parse;
mod error;
pub use error::*;
Expand All @@ -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
Expand All @@ -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 {
Expand Down

0 comments on commit 6137e91

Please sign in to comment.