use core::fmt;
use std::fmt::Formatter;
use std::sync::Arc;
use serde::{Deserialize, Serialize};
use typed_builder::TypedBuilder;
use super::transform::Transform;
use crate::error::Result;
use crate::spec::Schema;
use crate::{Error, ErrorKind};
pub type SortOrderRef = Arc<SortOrder>;
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Copy, Clone)]
pub enum SortDirection {
#[serde(rename = "asc")]
Ascending,
#[serde(rename = "desc")]
Descending,
}
impl fmt::Display for SortDirection {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match *self {
SortDirection::Ascending => write!(f, "ascending"),
SortDirection::Descending => write!(f, "descending"),
}
}
}
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Copy, Clone)]
pub enum NullOrder {
#[serde(rename = "nulls-first")]
First,
#[serde(rename = "nulls-last")]
Last,
}
impl fmt::Display for NullOrder {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match *self {
NullOrder::First => write!(f, "first"),
NullOrder::Last => write!(f, "last"),
}
}
}
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone, TypedBuilder)]
#[serde(rename_all = "kebab-case")]
pub struct SortField {
pub source_id: i32,
pub transform: Transform,
pub direction: SortDirection,
pub null_order: NullOrder,
}
impl fmt::Display for SortField {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(
f,
"SortField {{ source_id: {}, transform: {}, direction: {}, null_order: {} }}",
self.source_id, self.transform, self.direction, self.null_order
)
}
}
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone, Builder, Default)]
#[serde(rename_all = "kebab-case")]
#[builder(setter(prefix = "with"))]
#[builder(build_fn(skip))]
pub struct SortOrder {
#[builder(default)]
pub order_id: i64,
#[builder(setter(each(name = "with_sort_field")), default)]
pub fields: Vec<SortField>,
}
impl SortOrder {
pub(crate) const UNSORTED_ORDER_ID: i64 = 0;
pub fn builder() -> SortOrderBuilder {
SortOrderBuilder::default()
}
pub fn unsorted_order() -> SortOrder {
SortOrder {
order_id: SortOrder::UNSORTED_ORDER_ID,
fields: Vec::new(),
}
}
pub fn is_unsorted(&self) -> bool {
self.fields.is_empty()
}
pub fn with_order_id(self, order_id: i64) -> SortOrder {
SortOrder {
order_id,
fields: self.fields,
}
}
}
impl SortOrderBuilder {
pub fn build_unbound(&self) -> Result<SortOrder> {
let fields = self.fields.clone().unwrap_or_default();
return match (self.order_id, fields.as_slice()) {
(Some(SortOrder::UNSORTED_ORDER_ID) | None, []) => Ok(SortOrder::unsorted_order()),
(_, []) => Err(Error::new(
ErrorKind::Unexpected,
format!("Unsorted order ID must be {}", SortOrder::UNSORTED_ORDER_ID),
)),
(Some(SortOrder::UNSORTED_ORDER_ID), [..]) => Err(Error::new(
ErrorKind::Unexpected,
format!(
"Sort order ID {} is reserved for unsorted order",
SortOrder::UNSORTED_ORDER_ID
),
)),
(maybe_order_id, [..]) => Ok(SortOrder {
order_id: maybe_order_id.unwrap_or(1),
fields: fields.to_vec(),
}),
};
}
pub fn build(&self, schema: &Schema) -> Result<SortOrder> {
let unbound_sort_order = self.build_unbound()?;
SortOrderBuilder::check_compatibility(unbound_sort_order, schema)
}
fn check_compatibility(sort_order: SortOrder, schema: &Schema) -> Result<SortOrder> {
let sort_fields = &sort_order.fields;
for sort_field in sort_fields {
match schema.field_by_id(sort_field.source_id) {
None => {
return Err(Error::new(
ErrorKind::Unexpected,
format!("Cannot find source column for sort field: {sort_field}"),
))
}
Some(source_field) => {
let source_type = source_field.field_type.as_ref();
if !source_type.is_primitive() {
return Err(Error::new(
ErrorKind::Unexpected,
format!("Cannot sort by non-primitive source field: {source_type}"),
));
}
let field_transform = sort_field.transform;
if field_transform.result_type(source_type).is_err() {
return Err(Error::new(
ErrorKind::Unexpected,
format!(
"Invalid source type {source_type} for transform {field_transform}"
),
));
}
}
}
}
Ok(sort_order)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::spec::{ListType, NestedField, PrimitiveType, Type};
#[test]
fn test_sort_field() {
let spec = r#"
{
"transform": "bucket[4]",
"source-id": 3,
"direction": "desc",
"null-order": "nulls-last"
}
"#;
let field: SortField = serde_json::from_str(spec).unwrap();
assert_eq!(Transform::Bucket(4), field.transform);
assert_eq!(3, field.source_id);
assert_eq!(SortDirection::Descending, field.direction);
assert_eq!(NullOrder::Last, field.null_order);
}
#[test]
fn test_sort_order() {
let spec = r#"
{
"order-id": 1,
"fields": [ {
"transform": "identity",
"source-id": 2,
"direction": "asc",
"null-order": "nulls-first"
}, {
"transform": "bucket[4]",
"source-id": 3,
"direction": "desc",
"null-order": "nulls-last"
} ]
}
"#;
let order: SortOrder = serde_json::from_str(spec).unwrap();
assert_eq!(Transform::Identity, order.fields[0].transform);
assert_eq!(2, order.fields[0].source_id);
assert_eq!(SortDirection::Ascending, order.fields[0].direction);
assert_eq!(NullOrder::First, order.fields[0].null_order);
assert_eq!(Transform::Bucket(4), order.fields[1].transform);
assert_eq!(3, order.fields[1].source_id);
assert_eq!(SortDirection::Descending, order.fields[1].direction);
assert_eq!(NullOrder::Last, order.fields[1].null_order);
}
#[test]
fn test_build_unbound_should_return_err_if_unsorted_order_does_not_have_an_order_id_of_zero() {
assert_eq!(
SortOrder::builder()
.with_order_id(1)
.build_unbound()
.expect_err("Expected an Err value")
.message(),
"Unsorted order ID must be 0"
)
}
#[test]
fn test_build_unbound_should_return_err_if_order_id_equals_zero_is_used_for_anything_other_than_unsorted_order(
) {
assert_eq!(
SortOrder::builder()
.with_order_id(SortOrder::UNSORTED_ORDER_ID)
.with_sort_field(
SortField::builder()
.source_id(2)
.direction(SortDirection::Ascending)
.null_order(NullOrder::First)
.transform(Transform::Identity)
.build()
)
.build_unbound()
.expect_err("Expected an Err value")
.message(),
"Sort order ID 0 is reserved for unsorted order"
)
}
#[test]
fn test_build_unbound_returns_correct_default_order_id_for_no_fields() {
assert_eq!(
SortOrder::builder()
.build_unbound()
.expect("Expected an Ok value")
.order_id,
SortOrder::UNSORTED_ORDER_ID
)
}
#[test]
fn test_build_unbound_returns_correct_default_order_id_for_fields() {
let sort_field = SortField::builder()
.source_id(2)
.direction(SortDirection::Ascending)
.null_order(NullOrder::First)
.transform(Transform::Identity)
.build();
assert_ne!(
SortOrder::builder()
.with_sort_field(sort_field.clone())
.build_unbound()
.expect("Expected an Ok value")
.order_id,
SortOrder::UNSORTED_ORDER_ID
)
}
#[test]
fn test_build_unbound_should_return_unsorted_sort_order() {
assert_eq!(
SortOrder::builder()
.with_order_id(SortOrder::UNSORTED_ORDER_ID)
.build_unbound()
.expect("Expected an Ok value"),
SortOrder::unsorted_order()
)
}
#[test]
fn test_build_unbound_should_return_sort_order_with_given_order_id_and_sort_fields() {
let sort_field = SortField::builder()
.source_id(2)
.direction(SortDirection::Ascending)
.null_order(NullOrder::First)
.transform(Transform::Identity)
.build();
assert_eq!(
SortOrder::builder()
.with_order_id(2)
.with_sort_field(sort_field.clone())
.build_unbound()
.expect("Expected an Ok value"),
SortOrder {
order_id: 2,
fields: vec![sort_field]
}
)
}
#[test]
fn test_build_unbound_should_return_sort_order_with_given_sort_fields_and_defaults_to_1_if_missing_an_order_id(
) {
let sort_field = SortField::builder()
.source_id(2)
.direction(SortDirection::Ascending)
.null_order(NullOrder::First)
.transform(Transform::Identity)
.build();
assert_eq!(
SortOrder::builder()
.with_sort_field(sort_field.clone())
.build_unbound()
.expect("Expected an Ok value"),
SortOrder {
order_id: 1,
fields: vec![sort_field]
}
)
}
#[test]
fn test_build_should_return_err_if_sort_order_field_is_not_present_in_schema() {
let schema = Schema::builder()
.with_schema_id(1)
.with_fields(vec![NestedField::required(
1,
"foo",
Type::Primitive(PrimitiveType::Int),
)
.into()])
.build()
.unwrap();
let sort_order_builder_result = SortOrder::builder()
.with_sort_field(
SortField::builder()
.source_id(2)
.direction(SortDirection::Ascending)
.null_order(NullOrder::First)
.transform(Transform::Identity)
.build(),
)
.build(&schema);
assert_eq!(
sort_order_builder_result
.expect_err("Expected an Err value")
.message(),
"Cannot find source column for sort field: SortField { source_id: 2, transform: identity, direction: ascending, null_order: first }"
)
}
#[test]
fn test_build_should_return_err_if_source_field_is_not_a_primitive_type() {
let schema = Schema::builder()
.with_schema_id(1)
.with_fields(vec![NestedField::required(
1,
"foo",
Type::List(ListType {
element_field: NestedField::list_element(
2,
Type::Primitive(PrimitiveType::String),
true,
)
.into(),
}),
)
.into()])
.build()
.unwrap();
let sort_order_builder_result = SortOrder::builder()
.with_sort_field(
SortField::builder()
.source_id(1)
.direction(SortDirection::Ascending)
.null_order(NullOrder::First)
.transform(Transform::Identity)
.build(),
)
.build(&schema);
assert_eq!(
sort_order_builder_result
.expect_err("Expected an Err value")
.message(),
"Cannot sort by non-primitive source field: list"
)
}
#[test]
fn test_build_should_return_err_if_source_field_type_is_not_supported_by_transform() {
let schema = Schema::builder()
.with_schema_id(1)
.with_fields(vec![NestedField::required(
1,
"foo",
Type::Primitive(PrimitiveType::Int),
)
.into()])
.build()
.unwrap();
let sort_order_builder_result = SortOrder::builder()
.with_sort_field(
SortField::builder()
.source_id(1)
.direction(SortDirection::Ascending)
.null_order(NullOrder::First)
.transform(Transform::Year)
.build(),
)
.build(&schema);
assert_eq!(
sort_order_builder_result
.expect_err("Expected an Err value")
.message(),
"Invalid source type int for transform year"
)
}
#[test]
fn test_build_should_return_valid_sort_order() {
let schema = Schema::builder()
.with_schema_id(1)
.with_fields(vec![
NestedField::required(1, "foo", Type::Primitive(PrimitiveType::String)).into(),
NestedField::required(2, "bar", Type::Primitive(PrimitiveType::Int)).into(),
])
.build()
.unwrap();
let sort_field = SortField::builder()
.source_id(2)
.direction(SortDirection::Ascending)
.null_order(NullOrder::First)
.transform(Transform::Identity)
.build();
let sort_order_builder_result = SortOrder::builder()
.with_sort_field(sort_field.clone())
.build(&schema);
assert_eq!(
sort_order_builder_result.expect("Expected an Ok value"),
SortOrder {
order_id: 1,
fields: vec![sort_field],
}
)
}
}