Support :with (#311) (#414) r=nalexander

* Pre: refactor projector code.
* Part 1: maintain 'with' variables in AlgebrizedQuery.
* Part 2: include necessary 'with' variables in SQL projection list.

The test produces projection elements for `:with`, even though there are
no aggregates in the query. This test will need to be adjusted when we
optimize this away!

query-algebrizer/src/lib.rs

@@ -14,6 +14,8 @@ extern crate error_chain;
 extern crate mentat_core;
 extern crate mentat_query;
 
+use std::collections::BTreeSet;
+
 mod errors;
 mod types;
 mod validate;
@@ -29,6 +31,7 @@ use mentat_query::{
     FindQuery,
     FindSpec,
     SrcVar,
+    Variable,
 };
 
 pub use errors::{
@@ -41,6 +44,7 @@ pub use errors::{
 pub struct AlgebraicQuery {
     default_source: SrcVar,
     pub find_spec: FindSpec,
+    pub with: BTreeSet<Variable>,
     has_aggregates: bool,
     pub limit: Option<u64>,
     pub cc: clauses::ConjoiningClauses,
@@ -96,6 +100,7 @@ pub fn algebrize_with_cc(schema: &Schema, parsed: FindQuery, mut cc: ConjoiningC
         default_source: parsed.default_source,
         find_spec: parsed.find_spec,
         has_aggregates: false,           // TODO: we don't parse them yet.
+        with: parsed.with.into_iter().collect(),
         limit: limit,
         cc: cc,
     })

query-projector/src/lib.rs

@@ -37,6 +37,7 @@ use mentat_db::{
 use mentat_query::{
     Element,
     FindSpec,
+    Variable,
 };
 
 use mentat_query_algebrizer::{
@@ -47,6 +48,7 @@ use mentat_query_algebrizer::{
 
 use mentat_query_sql::{
     ColumnOrExpression,
+    Name,
     Projection,
     ProjectedColumn,
 };
@@ -154,6 +156,29 @@ impl TypedIndex {
     }
 }
 
+fn candidate_column(query: &AlgebraicQuery, var: &Variable) -> (ColumnOrExpression, Name) {
+    // Every variable should be bound by the top-level CC to at least
+    // one column in the query. If that constraint is violated it's a
+    // bug in our code, so it's appropriate to panic here.
+    let columns = query.cc
+                       .column_bindings
+                       .get(var)
+                       .expect("Every variable has a binding");
+
+    let qa = columns[0].clone();
+    let name = VariableColumn::Variable(var.clone()).column_name();
+    (ColumnOrExpression::Column(qa), name)
+}
+
+fn candidate_type_column(query: &AlgebraicQuery, var: &Variable) -> (ColumnOrExpression, Name) {
+    let extracted_alias = query.cc
+                               .extracted_types
+                               .get(var)
+                               .expect("Every variable has a known type or an extracted type");
+    let type_name = VariableColumn::VariableTypeTag(var.clone()).column_name();
+    (ColumnOrExpression::Column(extracted_alias.clone()), type_name)
+}
+
 /// Walk an iterator of `Element`s, collecting projector templates and columns.
 ///
 /// Returns a pair: the SQL projection (which should always be a `Projection::Columns`)
@@ -174,6 +199,7 @@ fn project_elements<'a, I: IntoIterator<Item = &'a Element>>(
     let mut cols = Vec::with_capacity(count);
     let mut i: i32 = 0;
     let mut templates = vec![];
+    let mut with = query.with.clone();
 
     for e in elements {
         match e {
@@ -181,39 +207,39 @@ fn project_elements<'a, I: IntoIterator<Item = &'a Element>>(
             // into the SQL projection, aliased to the name of the variable,
             // and we push an annotated index into the projector.
             &Element::Variable(ref var) => {
-                // Every variable should be bound by the top-level CC to at least
-                // one column in the query. If that constraint is violated it's a
-                // bug in our code, so it's appropriate to panic here.
-                let columns = query.cc
-                                   .column_bindings
-                                   .get(var)
-                                   .expect("Every variable has a binding");
-
-                let qa = columns[0].clone();
-                let name = VariableColumn::Variable(var.clone()).column_name();
+                // If we're projecting this, we don't need it in :with.
+                with.remove(var);
 
+                let (column, name) = candidate_column(query, var);
+                cols.push(ProjectedColumn(column, name));
                 if let Some(t) = query.cc.known_type(var) {
-                    cols.push(ProjectedColumn(ColumnOrExpression::Column(qa), name));
                     let tag = t.value_type_tag();
                     templates.push(TypedIndex::Known(i, tag));
                     i += 1;     // We used one SQL column.
                 } else {
-                    cols.push(ProjectedColumn(ColumnOrExpression::Column(qa), name));
                     templates.push(TypedIndex::Unknown(i, i + 1));
                     i += 2;     // We used two SQL columns.
 
                     // Also project the type from the SQL query.
-                    let extracted_alias = query.cc
-                                               .extracted_types
-                                               .get(var)
-                                               .expect("Every variable has a known type or an extracted type");
-                    let type_name = VariableColumn::VariableTypeTag(var.clone()).column_name();
-                    cols.push(ProjectedColumn(ColumnOrExpression::Column(extracted_alias.clone()), type_name));
+                    let (type_column, type_name) = candidate_type_column(query, &var);
+                    cols.push(ProjectedColumn(type_column, type_name));
                 }
             }
         }
     }
 
+    for var in with {
+        // We need to collect these into the SQL column list, but they don't affect projection.
+        // If a variable is of a non-fixed type, also project the type tag column, so we don't
+        // accidentally unify across types when considering uniqueness!
+        let (column, name) = candidate_column(query, &var);
+        cols.push(ProjectedColumn(column, name));
+        if query.cc.known_type(&var).is_none() {
+            let (type_column, type_name) = candidate_type_column(query, &var);
+            cols.push(ProjectedColumn(type_column, type_name));
+        }
+    }
+
     (Projection::Columns(cols), templates)
 }
 

query-translator/tests/translate.rs

@@ -353,3 +353,20 @@ fn test_complex_or_join_type_projection() {
                     LIMIT 1");
     assert_eq!(args, vec![]);
 }
+
+#[test]
+fn test_with_without_aggregate() {
+    let schema = prepopulated_schema();
+
+    // Known type.
+    let input = r#"[:find ?x :with ?y :where [?x :foo/bar ?y]]"#;
+    let SQLQuery { sql, args } = translate(&schema, input, None);
+    assert_eq!(sql, "SELECT DISTINCT `datoms00`.e AS `?x`, `datoms00`.v AS `?y` FROM `datoms` AS `datoms00` WHERE `datoms00`.a = 99");
+    assert_eq!(args, vec![]);
+
+    // Unknown type.
+    let input = r#"[:find ?x :with ?y :where [?x _ ?y]]"#;
+    let SQLQuery { sql, args } = translate(&schema, input, None);
+    assert_eq!(sql, "SELECT DISTINCT `all_datoms00`.e AS `?x`, `all_datoms00`.v AS `?y`, `all_datoms00`.value_type_tag AS `?y_value_type_tag` FROM `all_datoms` AS `all_datoms00`");
+    assert_eq!(args, vec![]);
+}