Allow ORDER BY to 'see' projection names (#443)

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add `partiql-extension-visualize` for visualizing AST and logical plan
 
 ### Fixed
+- Fixed `ORDER BY`'s ability to see into projection aliases
 
 ## [0.6.0] - 2023-10-31
 ### Changed

partiql-eval/src/eval/expr/path.rs

@@ -1,6 +1,6 @@
 use crate::env::Bindings;
 
-pub use core::borrow::{Borrow, BorrowMut};
+pub use core::borrow::Borrow;
 
 use crate::eval::expr::{BindError, BindEvalExpr, EvalExpr};
 use crate::eval::EvalContext;

partiql-logical-planner/src/lower.rs

@@ -32,10 +32,12 @@ use partiql_catalog::call_defs::{CallArgument, CallDef};
 
 use partiql_ast_passes::error::{AstTransformError, AstTransformationError};
 
+use partiql_ast_passes::name_resolver::NameRef;
 use partiql_catalog::Catalog;
 use partiql_extension_ion::decode::{IonDecoderBuilder, IonDecoderConfig};
 use partiql_extension_ion::Encoding;
 use partiql_logical::AggFunc::{AggAny, AggAvg, AggCount, AggEvery, AggMax, AggMin, AggSum};
+use partiql_logical::ValueExpr::DynamicLookup;
 use std::sync::atomic::{AtomicU32, Ordering};
 
 type FnvIndexMap<K, V> = IndexMap<K, V, FnvBuildHasher>;
@@ -161,6 +163,8 @@ pub struct AstToLogical<'a> {
 
     from_lets: HashSet<ast::NodeId>,
 
+    projection_renames: Vec<FnvIndexMap<String, BindingsName<'a>>>,
+
     aliases: FnvIndexMap<NodeId, SymbolPrimitive>,
 
     // generator of 'fresh' ids
@@ -230,6 +234,8 @@ impl<'a> AstToLogical<'a> {
 
             from_lets: Default::default(),
 
+            projection_renames: Default::default(),
+
             aliases: Default::default(),
 
             // generator of 'fresh' ids
@@ -283,6 +289,17 @@ impl<'a> AstToLogical<'a> {
     }
 
     fn resolve_varref(&self, varref: &ast::VarRef) -> logical::ValueExpr {
+        fn binding_to_static<'a>(binding: &'a BindingsName<'a>) -> BindingsName<'static> {
+            match binding {
+                BindingsName::CaseSensitive(n) => {
+                    BindingsName::CaseSensitive(Cow::Owned(n.as_ref().to_string()))
+                }
+                BindingsName::CaseInsensitive(n) => {
+                    BindingsName::CaseInsensitive(Cow::Owned(n.as_ref().to_string()))
+                }
+            }
+        }
+
         // Convert a `SymbolPrimitive` into a `BindingsName`
         fn symprim_to_binding(sym: &SymbolPrimitive) -> BindingsName<'static> {
             match sym.case {
@@ -306,14 +323,36 @@ impl<'a> AstToLogical<'a> {
             if let Some(key_schema) = self.key_registry.schema.get(id) {
                 let key_schema: &name_resolver::KeySchema = key_schema;
 
-                let name_ref: &name_resolver::NameRef = key_schema
+                let name_ref: &NameRef = key_schema
                     .consume
                     .iter()
                     .find(|name_ref| name_ref.sym == varref.name)
                     .expect("NameRef");
 
                 let var_binding = symprim_to_binding(&name_ref.sym);
                 let mut lookups = vec![];
+
+                if matches!(self.current_ctx(), Some(QueryContext::Order)) {
+                    if let Some(renames) = self.projection_renames.last() {
+                        let binding = renames
+                            .iter()
+                            .find(|(k, _)| {
+                                let SymbolPrimitive { value, case } = &name_ref.sym;
+                                match case {
+                                    CaseSensitivity::CaseSensitive => value == *k,
+                                    CaseSensitivity::CaseInsensitive => unicase::eq(value, *k),
+                                }
+                            })
+                            .map(|(_k, v)| binding_to_static(v))
+                            .unwrap_or_else(|| symprim_to_binding(&name_ref.sym));
+
+                        lookups.push(DynamicLookup(Box::new(vec![ValueExpr::VarRef(
+                            binding,
+                            VarRefType::Local,
+                        )])));
+                    }
+                }
+
                 for lookup in &name_ref.lookup {
                     match lookup {
                         name_resolver::NameLookup::Global => {
@@ -442,11 +481,13 @@ impl<'a> AstToLogical<'a> {
     fn enter_q(&mut self) {
         self.q_stack.push(Default::default());
         self.ctx_stack.push(QueryContext::Query);
+        self.projection_renames.push(Default::default());
     }
 
     #[inline]
     fn exit_q(&mut self) -> QueryClauses {
-        self.ctx_stack.pop();
+        self.projection_renames.pop().expect("q level");
+        self.ctx_stack.pop().expect("q level");
         self.q_stack.pop().expect("q level")
     }
 
@@ -856,6 +897,15 @@ impl<'a, 'ast> Visitor<'ast> for AstToLogical<'a> {
                             "".to_string()
                         }
                     };
+
+                    if !alias.is_empty() {
+                        if let ValueExpr::VarRef(name, _vrtype) = &value {
+                            self.projection_renames
+                                .last_mut()
+                                .expect("renames")
+                                .insert(alias.clone(), name.clone());
+                        }
+                    }
                     exprs.push((alias, value));
                 }
 
@@ -1286,16 +1336,16 @@ impl<'a, 'ast> Visitor<'ast> for AstToLogical<'a> {
         Traverse::Continue
     }
 
-    fn enter_var_ref(&mut self, _var_ref: &'ast VarRef) -> Traverse {
+    fn enter_var_ref(&mut self, var_ref: &'ast VarRef) -> Traverse {
         let is_path = matches!(self.current_ctx(), Some(QueryContext::Path));
         if !is_path {
-            let options = self.resolve_varref(_var_ref);
+            let options = self.resolve_varref(var_ref);
             self.push_vexpr(options);
         } else {
             let VarRef {
                 name: SymbolPrimitive { value, case },
                 qualifier: _,
-            } = _var_ref;
+            } = var_ref;
             let name = match case {
                 CaseSensitivity::CaseSensitive => {
                     BindingsName::CaseSensitive(Cow::Owned(value.clone()))
@@ -1963,7 +2013,7 @@ mod tests {
         let lowering_errs = logical.expect_err("Expect errs").errors;
         assert_eq!(lowering_errs.len(), 2);
         assert_eq!(
-            lowering_errs.get(0),
+            lowering_errs.first(),
             Some(&AstTransformError::UnsupportedFunction("foo".to_string()))
         );
         assert_eq!(
@@ -1985,7 +2035,7 @@ mod tests {
         let lowering_errs = logical.expect_err("Expect errs").errors;
         assert_eq!(lowering_errs.len(), 2);
         assert_eq!(
-            lowering_errs.get(0),
+            lowering_errs.first(),
             Some(&AstTransformError::InvalidNumberOfArguments(
                 "abs".to_string()
             ))

partiql/Cargo.toml

@@ -25,14 +25,18 @@ bench = false
 [dev-dependencies]
 partiql-parser = { path = "../partiql-parser" }
 partiql-ast = { path = "../partiql-ast" }
+partiql-ast-passes = { path = "../partiql-ast-passes" }
 partiql-catalog = { path = "../partiql-catalog"}
 partiql-value = { path = "../partiql-value" }
 partiql-logical = { path = "../partiql-logical" }
 partiql-logical-planner = { path = "../partiql-logical-planner" }
 partiql-eval = { path = "../partiql-eval" }
 
-itertools = "0.10"
-criterion = "0.4"
+
+thiserror = "1.0"
+
+itertools = "0.12"
+criterion = "0.5"
 rand = "0.8"
 
 [[bench]]

partiql/src/lib.rs

@@ -1,5 +1,112 @@
 #[cfg(test)]
 mod tests {
+    use partiql_ast_passes::error::AstTransformationError;
+    use partiql_catalog::{Catalog, PartiqlCatalog};
+    use partiql_eval as eval;
+    use partiql_eval::env::basic::MapBindings;
+    use partiql_eval::error::{EvalErr, PlanErr};
+    use partiql_eval::eval::{EvalPlan, EvalResult, Evaluated};
+    use partiql_eval::plan::EvaluationMode;
+    use partiql_logical as logical;
+    use partiql_parser::{Parsed, ParserError, ParserResult};
+    use partiql_value::Value;
+    use thiserror::Error;
+
+    #[derive(Error, Debug)]
+    enum TestError<'a> {
+        #[error("Parse error: {0:?}")]
+        Parse(ParserError<'a>),
+        #[error("Lower error: {0:?}")]
+        Lower(AstTransformationError),
+        #[error("Plan error: {0:?}")]
+        Plan(PlanErr),
+        #[error("Evaluation error: {0:?}")]
+        Eval(EvalErr),
+    }
+
+    impl<'a> From<ParserError<'a>> for TestError<'a> {
+        fn from(err: ParserError<'a>) -> Self {
+            TestError::Parse(err)
+        }
+    }
+
+    impl From<AstTransformationError> for TestError<'_> {
+        fn from(err: AstTransformationError) -> Self {
+            TestError::Lower(err)
+        }
+    }
+
+    impl From<PlanErr> for TestError<'_> {
+        fn from(err: PlanErr) -> Self {
+            TestError::Plan(err)
+        }
+    }
+
+    impl From<EvalErr> for TestError<'_> {
+        fn from(err: EvalErr) -> Self {
+            TestError::Eval(err)
+        }
+    }
+
+    #[track_caller]
+    #[inline]
+    fn parse(statement: &str) -> ParserResult {
+        partiql_parser::Parser::default().parse(statement)
+    }
+
+    #[track_caller]
+    #[inline]
+    fn lower(
+        catalog: &dyn Catalog,
+        parsed: &Parsed,
+    ) -> Result<logical::LogicalPlan<logical::BindingsOp>, AstTransformationError> {
+        let planner = partiql_logical_planner::LogicalPlanner::new(catalog);
+        planner.lower(parsed)
+    }
+
+    #[track_caller]
+    #[inline]
+    fn compile(
+        mode: EvaluationMode,
+        catalog: &dyn Catalog,
+        logical: logical::LogicalPlan<logical::BindingsOp>,
+    ) -> Result<EvalPlan, PlanErr> {
+        let mut planner = eval::plan::EvaluatorPlanner::new(mode, catalog);
+        planner.compile(&logical)
+    }
+
+    #[track_caller]
+    #[inline]
+    fn evaluate(mut plan: EvalPlan, bindings: MapBindings<Value>) -> EvalResult {
+        plan.execute_mut(bindings)
+    }
+
+    #[track_caller]
+    #[inline]
+    fn eval(statement: &str, mode: EvaluationMode) -> Result<Evaluated, TestError<'_>> {
+        let catalog = PartiqlCatalog::default();
+
+        let parsed = parse(statement)?;
+        let lowered = lower(&catalog, &parsed)?;
+        let bindings = Default::default();
+        let plan = compile(mode, &catalog, lowered)?;
+        Ok(evaluate(plan, bindings)?)
+    }
+
     #[test]
-    fn todo() {}
+    fn order_by_count() {
+        let query = "select foo, count(1) as n from
+            <<
+            { 'foo': 'foo' },
+        { 'foo': 'bar' },
+        { 'foo': 'qux' },
+        { 'foo': 'bar' },
+        { 'foo': 'baz' },
+        { 'foo': 'bar' },
+        { 'foo': 'baz' }
+        >>  group by foo order by n desc";
+
+        let res = eval(query, EvaluationMode::Permissive);
+        assert!(res.is_ok());
+    }
 }