Implements ORDER BY (#337)

jpschorr · web-flow · commit 11e5623b6892 · 2023-04-10T10:12:56.000-07:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -18,6 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Expose `partiql_value::parse_ion` as a public API.
 - Implements `GROUP BY` operator in evaluator
 - Implements `HAVING` operator in evaluator
+- Implements `ORDER BY` operator in evaluator
 
 ### Fixes
 - Fixes Tuple value duplicate equality and hashing
diff --git a/partiql-eval/src/eval/evaluable.rs b/partiql-eval/src/eval/evaluable.rs
@@ -6,6 +6,7 @@ use partiql_value::Value::{Boolean, Missing, Null};
 use partiql_value::{partiql_bag, partiql_tuple, Bag, List, Tuple, Value};
 use std::borrow::{Borrow, Cow};
 use std::cell::RefCell;
+use std::cmp::Ordering;
 use std::collections::HashMap;
 use std::fmt::Debug;
 use std::rc::Rc;
@@ -548,6 +549,84 @@ impl Evaluable for EvalHaving {
     }
 }
 
+#[derive(Debug)]
+pub struct EvalOrderBySortCondition {
+    pub expr: Box<dyn EvalExpr>,
+    pub spec: EvalOrderBySortSpec,
+}
+
+#[derive(Debug)]
+pub enum EvalOrderBySortSpec {
+    AscNullsFirst,
+    AscNullsLast,
+    DescNullsFirst,
+    DescNullsLast,
+}
+
+/// Represents an evaluation `Order By` operator; e.g. `ORDER BY a DESC NULLS LAST` in `SELECT a FROM t ORDER BY a DESC NULLS LAST`.
+#[derive(Debug)]
+pub struct EvalOrderBy {
+    pub cmp: Vec<EvalOrderBySortCondition>,
+    pub input: Option<Value>,
+}
+
+impl EvalOrderBy {
+    #[inline]
+    fn compare(&self, l: &Value, r: &Value, ctx: &dyn EvalContext) -> Ordering {
+        let l = l.as_tuple_ref();
+        let r = r.as_tuple_ref();
+        self.cmp
+            .iter()
+            .map(|spec| {
+                let l = spec.expr.evaluate(&l, ctx);
+                let r = spec.expr.evaluate(&r, ctx);
+
+                match spec.spec {
+                    EvalOrderBySortSpec::AscNullsFirst => l.as_ref().cmp(r.as_ref()),
+                    EvalOrderBySortSpec::AscNullsLast => match (l.as_ref(), r.as_ref()) {
+                        (Null, Null) => Ordering::Equal,
+                        (Null, Missing) => Ordering::Less,
+                        (Missing, Missing) => Ordering::Equal,
+                        (Missing, Null) => Ordering::Greater,
+                        (Null, _) => Ordering::Greater,
+                        (Missing, _) => Ordering::Greater,
+                        (_, Null) => Ordering::Less,
+                        (_, Missing) => Ordering::Less,
+                        (l, r) => l.cmp(r),
+                    },
+                    EvalOrderBySortSpec::DescNullsFirst => match (l.as_ref(), r.as_ref()) {
+                        (Null, Null) => Ordering::Equal,
+                        (Null, Missing) => Ordering::Less,
+                        (Missing, Missing) => Ordering::Equal,
+                        (Missing, Null) => Ordering::Greater,
+                        (Null, _) => Ordering::Less,
+                        (Missing, _) => Ordering::Less,
+                        (_, Null) => Ordering::Greater,
+                        (_, Missing) => Ordering::Greater,
+                        (l, r) => r.cmp(l),
+                    },
+                    EvalOrderBySortSpec::DescNullsLast => r.as_ref().cmp(l.as_ref()),
+                }
+            })
+            .find_or_last(|o| o != &Ordering::Equal)
+            .unwrap_or(Ordering::Equal)
+    }
+}
+
+impl Evaluable for EvalOrderBy {
+    fn evaluate(&mut self, ctx: &dyn EvalContext) -> Option<Value> {
+        let input_value = self.input.take().expect("Error in retrieving input value");
+
+        let mut values = input_value.into_iter().collect_vec();
+        values.sort_by(|l, r| self.compare(l, r, ctx));
+        Some(Value::from(List::from(values)))
+    }
+
+    fn update_input(&mut self, input: Value, _branch_num: u8) {
+        self.input = Some(input);
+    }
+}
+
 /// Represents an evaluation `LIMIT` and/or `OFFSET` operator.
 #[derive(Debug)]
 pub struct EvalLimitOffset {
diff --git a/partiql-eval/src/plan.rs b/partiql-eval/src/plan.rs
@@ -6,11 +6,15 @@ use partiql_logical as logical;
 
 use partiql_logical::{
     BinaryOp, BindingsOp, CallName, GroupingStrategy, IsTypeExpr, JoinKind, LogicalPlan, OpId,
-    PathComponent, Pattern, PatternMatchExpr, SearchedCase, Type, UnaryOp, ValueExpr,
+    PathComponent, Pattern, PatternMatchExpr, SearchedCase, SortSpecNullOrder, SortSpecOrder, Type,
+    UnaryOp, ValueExpr,
 };
 
 use crate::eval;
-use crate::eval::evaluable::{EvalGroupingStrategy, EvalJoinKind, EvalSubQueryExpr, Evaluable};
+use crate::eval::evaluable::{
+    EvalGroupingStrategy, EvalJoinKind, EvalOrderBy, EvalOrderBySortCondition, EvalOrderBySortSpec,
+    EvalSubQueryExpr, Evaluable,
+};
 use crate::eval::expr::pattern_match::like_to_re_pattern;
 use crate::eval::expr::{
     EvalBagExpr, EvalBetweenExpr, EvalBinOp, EvalBinOpExpr, EvalDynamicLookup, EvalExpr, EvalFnAbs,
@@ -158,7 +162,30 @@ impl EvaluatorPlanner {
                 let expr = self.plan_values(expr);
                 Box::new(eval::evaluable::EvalExprQuery::new(expr))
             }
-            BindingsOp::OrderBy => todo!("OrderBy"),
+            BindingsOp::OrderBy(logical::OrderBy { specs }) => {
+                let cmp = specs
+                    .iter()
+                    .map(|spec| {
+                        let expr = self.plan_values(&spec.expr);
+                        let spec = match (&spec.order, &spec.null_order) {
+                            (SortSpecOrder::Asc, SortSpecNullOrder::First) => {
+                                EvalOrderBySortSpec::AscNullsFirst
+                            }
+                            (SortSpecOrder::Asc, SortSpecNullOrder::Last) => {
+                                EvalOrderBySortSpec::AscNullsLast
+                            }
+                            (SortSpecOrder::Desc, SortSpecNullOrder::First) => {
+                                EvalOrderBySortSpec::DescNullsFirst
+                            }
+                            (SortSpecOrder::Desc, SortSpecNullOrder::Last) => {
+                                EvalOrderBySortSpec::DescNullsFirst
+                            }
+                        };
+                        EvalOrderBySortCondition { expr, spec }
+                    })
+                    .collect_vec();
+                Box::new(EvalOrderBy { cmp, input: None })
+            }
             BindingsOp::LimitOffset(logical::LimitOffset { limit, offset }) => {
                 Box::new(eval::evaluable::EvalLimitOffset {
                     limit: limit.as_ref().map(|e| self.plan_values(e)),
diff --git a/partiql-logical-planner/src/lower.rs b/partiql-logical-planner/src/lower.rs
@@ -7,16 +7,17 @@ use partiql_ast::ast::{
     Assignment, Bag, Between, BinOp, BinOpKind, Call, CallAgg, CallArg, CallArgNamed,
     CaseSensitivity, CreateIndex, CreateTable, Ddl, DdlOp, Delete, Dml, DmlOp, DropIndex,
     DropTable, FromClause, FromLet, FromLetKind, GroupByExpr, GroupKey, GroupingStrategy, Insert,
-    InsertValue, Item, Join, JoinKind, JoinSpec, Like, List, Lit, NodeId, OnConflict, OrderByExpr,
-    Path, PathStep, ProjectExpr, Projection, ProjectionKind, Query, QuerySet, Remove, SearchedCase,
-    Select, Set, SetExpr, SetQuantifier, Sexp, SimpleCase, Struct, SymbolPrimitive, UniOp,
-    UniOpKind, VarRef,
+    InsertValue, Item, Join, JoinKind, JoinSpec, Like, List, Lit, NodeId, NullOrderingSpec,
+    OnConflict, OrderByExpr, OrderingSpec, Path, PathStep, ProjectExpr, Projection, ProjectionKind,
+    Query, QuerySet, Remove, SearchedCase, Select, Set, SetExpr, SetQuantifier, Sexp, SimpleCase,
+    SortSpec, Struct, SymbolPrimitive, UniOp, UniOpKind, VarRef,
 };
 use partiql_ast::visit::{Visit, Visitor};
 use partiql_logical as logical;
 use partiql_logical::{
     BagExpr, BetweenExpr, BindingsOp, IsTypeExpr, LikeMatch, LikeNonStringNonLiteralMatch,
-    ListExpr, LogicalPlan, OpId, PathComponent, Pattern, PatternMatchExpr, TupleExpr, ValueExpr,
+    ListExpr, LogicalPlan, OpId, PathComponent, Pattern, PatternMatchExpr, SortSpecOrder,
+    TupleExpr, ValueExpr,
 };
 
 use partiql_value::{BindingsName, Value};
@@ -26,6 +27,7 @@ use std::collections::{HashMap, HashSet};
 use crate::call_defs::{CallArgument, FnSymTab, FN_SYM_TAB};
 use crate::name_resolver;
 use itertools::Itertools;
+
 use std::sync::atomic::{AtomicU32, Ordering};
 
 type FnvIndexMap<K, V> = IndexMap<K, V, FnvBuildHasher>;
@@ -34,6 +36,7 @@ type FnvIndexMap<K, V> = IndexMap<K, V, FnvBuildHasher>;
 enum QueryContext {
     FromLet,
     Path,
+    Order,
     Query,
 }
 
@@ -104,6 +107,7 @@ pub struct AstToLogical {
     vexpr_stack: Vec<Vec<ValueExpr>>,
     arg_stack: Vec<Vec<CallArgument>>,
     path_stack: Vec<Vec<PathComponent>>,
+    sort_stack: Vec<Vec<logical::SortSpec>>,
 
     from_lets: HashSet<ast::NodeId>,
 
@@ -166,6 +170,7 @@ impl AstToLogical {
             vexpr_stack: Default::default(),
             arg_stack: Default::default(),
             path_stack: Default::default(),
+            sort_stack: Default::default(),
 
             from_lets: Default::default(),
 
@@ -406,6 +411,23 @@ impl AstToLogical {
     fn push_path_step(&mut self, step: PathComponent) {
         self.path_stack.last_mut().unwrap().push(step);
     }
+
+    #[inline]
+    fn enter_sort(&mut self) {
+        self.sort_stack.push(vec![]);
+        self.ctx_stack.push(QueryContext::Order);
+    }
+
+    #[inline]
+    fn exit_sort(&mut self) -> Vec<logical::SortSpec> {
+        self.ctx_stack.pop();
+        self.sort_stack.pop().expect("sort specs")
+    }
+
+    #[inline]
+    fn push_sort_spec(&mut self, spec: logical::SortSpec) {
+        self.sort_stack.last_mut().unwrap().push(spec);
+    }
 }
 
 // SQL (and therefore PartiQL) text (and therefore AST) is not lexically-scoped as is the
@@ -1158,12 +1180,48 @@ impl<'ast> Visitor<'ast> for AstToLogical {
     }
 
     fn enter_order_by_expr(&mut self, _order_by_expr: &'ast OrderByExpr) {
-        self.enter_env();
+        self.enter_sort();
     }
 
     fn exit_order_by_expr(&mut self, _order_by_expr: &'ast OrderByExpr) {
-        let _env = self.exit_env();
-        todo!("order by clause");
+        let specs = self.exit_sort();
+        let order_by = logical::BindingsOp::OrderBy(logical::OrderBy { specs });
+        let id = self.plan.add_operator(order_by);
+        self.current_clauses_mut().order_by_clause.replace(id);
+    }
+
+    fn enter_sort_spec(&mut self, _sort_spec: &'ast SortSpec) {
+        self.enter_env();
+    }
+
+    fn exit_sort_spec(&mut self, sort_spec: &'ast SortSpec) {
+        let mut env = self.exit_env();
+        assert_eq!(env.len(), 1);
+
+        let expr = env.pop().unwrap();
+        let order = match sort_spec
+            .ordering_spec
+            .as_ref()
+            .unwrap_or(&OrderingSpec::Asc)
+        {
+            OrderingSpec::Asc => logical::SortSpecOrder::Asc,
+            OrderingSpec::Desc => logical::SortSpecOrder::Desc,
+        };
+
+        let null_order = match sort_spec.null_ordering_spec {
+            None => match order {
+                SortSpecOrder::Asc => logical::SortSpecNullOrder::Last,
+                SortSpecOrder::Desc => logical::SortSpecNullOrder::First,
+            },
+            Some(NullOrderingSpec::First) => logical::SortSpecNullOrder::First,
+            Some(NullOrderingSpec::Last) => logical::SortSpecNullOrder::Last,
+        };
+
+        self.push_sort_spec(logical::SortSpec {
+            expr,
+            order,
+            null_order,
+        });
     }
 
     fn enter_limit_offset_clause(&mut self, _limit_offset: &'ast ast::LimitOffsetClause) {
diff --git a/partiql-logical/src/lib.rs b/partiql-logical/src/lib.rs
@@ -108,10 +108,10 @@ where
     /// Extends the logical plan with the given data flows.
     /// #Examples:
     /// ```
-    /// use partiql_logical::{BindingsOp, GroupBy, GroupingStrategy, LimitOffset, LogicalPlan};
+    /// use partiql_logical::{BindingsOp, GroupBy, GroupingStrategy, LimitOffset, LogicalPlan, OrderBy};
     /// let mut p: LogicalPlan<BindingsOp> = LogicalPlan::new();
     ///
-    /// let a = p.add_operator(BindingsOp::OrderBy);
+    /// let a = p.add_operator(BindingsOp::OrderBy(OrderBy{specs: vec![]}));
     /// let b = p.add_operator(BindingsOp::Sink);
     /// let c = p.add_operator(BindingsOp::LimitOffset(LimitOffset{limit:None, offset:None}));
     /// let d = p.add_operator(BindingsOp::GroupBy(GroupBy {
@@ -191,7 +191,7 @@ pub enum BindingsOp {
     Pivot(Pivot),
     Unpivot(Unpivot),
     Filter(Filter),
-    OrderBy,
+    OrderBy(OrderBy),
     LimitOffset(LimitOffset),
     Join(Join),
     SetOp,
@@ -249,6 +249,36 @@ pub struct Having {
     pub expr: ValueExpr,
 }
 
+/// [`OrderBy`] represents a sort operatyion, e.g. `ORDER BY a DESC NULLS LAST` in `SELECT a FROM t ORDER BY a DESC NULLS LAST`.
+#[derive(Debug, Clone, Eq, PartialEq)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub struct OrderBy {
+    pub specs: Vec<SortSpec>,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub enum SortSpecOrder {
+    Asc,
+    Desc,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub enum SortSpecNullOrder {
+    First,
+    Last,
+}
+
+/// Represents a PartiQL sort specification.
+#[derive(Debug, Clone, Eq, PartialEq)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub struct SortSpec {
+    pub expr: ValueExpr,
+    pub order: SortSpecOrder,
+    pub null_order: SortSpecNullOrder,
+}
+
 /// [`LimitOffset`] represents a possible limit and/or offset operator, e.g. `LIMIT 10 OFFSET 5` in `SELECT a FROM t LIMIT 10 OFFSET 5`.
 #[derive(Debug, Clone, Eq, PartialEq)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
@@ -597,7 +627,7 @@ mod tests {
     #[test]
     fn test_plan() {
         let mut p: LogicalPlan<BindingsOp> = LogicalPlan::new();
-        let a = p.add_operator(BindingsOp::OrderBy);
+        let a = p.add_operator(BindingsOp::OrderBy(OrderBy { specs: vec![] }));
         let b = p.add_operator(BindingsOp::Sink);
         let c = p.add_operator(BindingsOp::LimitOffset(LimitOffset {
             limit: None,
