Implement LENGTH function

cswinter · Dec 20, 2018 · 38ba86d · 38ba86d
1 parent 745c38a
commit 38ba86d
Show file tree

Hide file tree

Showing 6 changed files with 52 additions and 1 deletion.
diff --git a/src/engine/operators/functions.rs b/src/engine/operators/functions.rs
@@ -35,3 +35,11 @@ impl<'a> MapOp<&'a str, u8> for RegexMatch {
     }
     fn name() -> &'static str { "not" }
 }
+
+
+pub struct Length;
+
+impl<'a> MapOp<&'a str, i64> for Length {
+    fn apply(&self, s: &'a str) -> i64 { s.len() as i64 }
+    fn name() -> &'static str { "length" }
+}
diff --git a/src/engine/operators/vector_operator.rs b/src/engine/operators/vector_operator.rs
@@ -614,6 +614,10 @@ impl<'a> VecOperator<'a> {
         Box::new(MapOperator { input, output, map: RegexMatch { r: regex::Regex::new(r).unwrap() } })
     }
 
+    pub fn length(input: BufferRef<&'a str>, output: BufferRef<i64>) -> BoxedOperator<'a> {
+        Box::new(MapOperator { input, output, map: Length })
+    }
+
     pub fn aggregate(input: TypedBufferRef,
                      grouping: TypedBufferRef,
                      max_index: BufferRef<Scalar<i64>>,

diff --git a/src/engine/planning/query_plan.rs b/src/engine/planning/query_plan.rs
@@ -339,6 +339,11 @@ pub enum QueryPlan {
         #[output]
         matches: BufferRef<u8>,
     },
+    Length {
+        string: BufferRef<&'static str>,
+        #[output]
+        length: BufferRef<i64>,
+    },
     /// Outputs a vector of indices from `0..plan.len()`
     Indices {
         plan: TypedBufferRef,
@@ -806,6 +811,16 @@ impl QueryPlan {
                         }
                         planner.to_year(decoded).into()
                     }
+                    Func1Type::Length => {
+                        let decoded = match t.codec.clone() {
+                            Some(codec) => codec.decode(plan, planner),
+                            None => plan,
+                        };
+                        if t.decoded != BasicType::String {
+                            bail!(QueryError::TypeError, "Found length({:?}), expected length(string)", &t)
+                        }
+                        planner.length(decoded.str()?).into()
+                    }
                     Func1Type::Not => {
                         let decoded = match t.codec.clone() {
                             Some(codec) => codec.decode(plan, planner),
@@ -867,7 +882,7 @@ fn encoding_range(plan: &TypedBufferRef, qp: &QueryPlanner) -> Option<(i64, i64)
         LZ4Decode { bytes, .. } => encoding_range(&bytes.into(), qp),
         DeltaDecode { ref plan, .. } => encoding_range(plan, qp),
         AssembleNullable { ref data, .. } => encoding_range(data, qp),
-        UnpackStrings { .. } | UnhexpackStrings { .. } => None,
+        UnpackStrings { .. } | UnhexpackStrings { .. } | Length { ..  } => None,
         ref plan => {
             // TODO(clemens): many more cases where we can determine range
             error!("encoding_range not implement for {:?}", plan);
@@ -1151,6 +1166,7 @@ pub fn prepare<'a>(plan: QueryPlan, constant_vecs: &mut Vec<BoxedData<'a>>, resu
         QueryPlan::Not { input, not } => VecOperator::not(input, not),
         QueryPlan::ToYear { timestamp, year } => VecOperator::to_year(timestamp.i64()?, year.i64()?),
         QueryPlan::Regex { plan, regex, matches } => VecOperator::regex(plan, &regex, matches),
+        QueryPlan::Length { string, length } => VecOperator::length(string, length),
         QueryPlan::Indices { plan, indices } => VecOperator::indices(plan, indices),
         QueryPlan::SortBy { ranking, indices, desc, stable, permutation } => VecOperator::sort_by(ranking, indices, desc, stable, permutation)?,
         QueryPlan::TopN { ranking, n, desc, tmp_keys, top_n } => VecOperator::top_n(ranking, tmp_keys, n, desc, top_n)?,

diff --git a/src/syntax/expression.rs b/src/syntax/expression.rs
@@ -38,6 +38,7 @@ pub enum Func1Type {
     Not,
     IsNull,
     IsNotNull,
+    Length,
 }
 
 impl Expr {

diff --git a/src/syntax/parser.rs b/src/syntax/parser.rs
@@ -119,6 +119,13 @@ fn expr(node: &ASTNode) -> Result<Box<Expr>, QueryError> {
                 }
                 Expr::Func2(Func2Type::RegexMatch, expr(&args[0])?, expr(&args[1])?)
             }
+            "LENGTH" => {
+                if args.len() != 1 {
+                    return Err(QueryError::ParseError(
+                        "Expected one arguments in length function".to_string()));
+                }
+                Expr::Func1(Func1Type::Length, expr(&args[0])?)
+            }
             "COUNT" => {
                 if args.len() != 1 {
                     return Err(QueryError::ParseError(

diff --git a/tests/query_tests.rs b/tests/query_tests.rs
@@ -527,6 +527,21 @@ fn test_group_by_trip_id() {
     )
 }
 
+#[test]
+fn test_string_length() {
+    test_query_nyc(
+        "SELECT length(pickup_ntaname), pickup_ntaname, COUNT(0)
+         FROM default
+         ORDER BY length(pickup_ntaname) DESC
+         LIMIT 3;",
+        &[
+            vec![Int(56), Str("Todt Hill-Emerson Hill-Heartland Village-Lighthouse Hill"), Int(1)],
+            vec![Int(50), Str("Mariner\'s Harbor-Arlington-Port Ivory-Graniteville"), Int(3)],
+            vec![Int(48), Str("DUMBO-Vinegar Hill-Downtown Brooklyn-Boerum Hill"), Int(245)],
+        ],
+    )
+}
+
 #[test]
 fn test_group_by_negative_expression() {
     test_query_ec(