mtrlz/sql: support abs function

Author: benesch

src/materialize/dataflow/func.rs

@@ -223,6 +223,34 @@ pub fn neg_float64(a: Datum) -> Datum {
     Datum::from(-a.unwrap_float64())
 }
 
+pub fn abs_int32(a: Datum) -> Datum {
+    if a.is_null() {
+        return Datum::Null;
+    }
+    Datum::from(a.unwrap_int32().abs())
+}
+
+pub fn abs_int64(a: Datum) -> Datum {
+    if a.is_null() {
+        return Datum::Null;
+    }
+    Datum::from(a.unwrap_int64().abs())
+}
+
+pub fn abs_float32(a: Datum) -> Datum {
+    if a.is_null() {
+        return Datum::Null;
+    }
+    Datum::from(a.unwrap_float32().abs())
+}
+
+pub fn abs_float64(a: Datum) -> Datum {
+    if a.is_null() {
+        return Datum::Null;
+    }
+    Datum::from(a.unwrap_float64().abs())
+}
+
 pub fn eq(a: Datum, b: Datum) -> Datum {
     if a.is_null() || b.is_null() {
         return Datum::Null;
@@ -344,6 +372,10 @@ pub enum UnaryFunc {
     NegInt64,
     NegFloat32,
     NegFloat64,
+    AbsInt32,
+    AbsInt64,
+    AbsFloat32,
+    AbsFloat64,
 }
 
 impl UnaryFunc {
@@ -355,6 +387,10 @@ impl UnaryFunc {
             UnaryFunc::NegInt64 => neg_int64,
             UnaryFunc::NegFloat32 => neg_float32,
             UnaryFunc::NegFloat64 => neg_float64,
+            UnaryFunc::AbsInt32 => abs_int32,
+            UnaryFunc::AbsInt64 => abs_int64,
+            UnaryFunc::AbsFloat32 => abs_float32,
+            UnaryFunc::AbsFloat64 => abs_float64,
         }
     }
 }

src/materialize/interchange/avro.rs

@@ -12,7 +12,7 @@ use std::collections::HashMap;
 use url::Url;
 
 use crate::repr::{Datum, FType, Type};
-use ore::vec::VecExt;
+use ore::collections::CollectionExt;
 
 /// Converts an Apache Avro schema into a [`repr::Type`].
 pub fn parse_schema(schema: &str) -> Result<Type, Error> {

src/materialize/sql/mod.rs

@@ -30,7 +30,7 @@ use crate::dataflow::{
 use crate::glue::*;
 use crate::interchange::avro;
 use crate::repr::{Datum, FType, Type};
-use ore::vec::VecExt;
+use ore::collections::CollectionExt;
 use plan::SQLPlan;
 use store::{DataflowStore, RemoveMode};
 
@@ -1086,22 +1086,42 @@ impl Planner {
         &self,
         ctx: &ExprContext,
         name: &'a SQLObjectName,
-        _args: &'a [ASTNode],
+        args: &'a [ASTNode],
         _over: &'a Option<SQLWindowSpec>,
         _all: bool,
         _distinct: bool,
         plan: &SQLPlan,
     ) -> Result<(Expr, Type), failure::Error> {
         let ident = name.to_string().to_lowercase();
+
         if AggregateFunc::is_aggregate_func(&ident) {
             if !ctx.allow_aggregates {
                 bail!("aggregate functions are not allowed in {}", ctx.scope);
             }
             let (i, typ) = plan.resolve_func(name);
             let expr = Expr::Column(i, Box::new(Expr::Ambient));
-            Ok((expr, typ.clone()))
-        } else {
-            bail!("Unsupported function: {}", ident)
+            return Ok((expr, typ.clone()))
+        }
+
+        match ident.as_str() {
+            "abs" => {
+                if args.len() != 1 {
+                    bail!("abs expects one argument, got {}", args.len());
+                }
+                let (expr, typ) = self.plan_expr(ctx, args.into_element(), plan)?;
+                let func = match typ.ftype {
+                    FType::Int32 => UnaryFunc::AbsInt32,
+                    FType::Int64 => UnaryFunc::AbsInt64,
+                    FType::Float32 => UnaryFunc::AbsFloat32,
+                    FType::Float64 => UnaryFunc::AbsFloat64,
+                    _ => bail!("abs does not accept arguments of type {:?}", typ),
+                };
+                Ok((Expr::CallUnary {
+                    func,
+                    expr: Box::new(expr),
+                }, typ))
+            }
+            _ => bail!("unsupported function: {}", ident),
         }
     }
 

src/ore/collections.rs

@@ -0,0 +1,33 @@
+// Copyright 2019 Materialize, Inc. All rights reserved.
+//
+// This file is part of Materialize. Materialize may not be used or
+// distributed without the express permission of Materialize, Inc.
+
+//! Collection utilities.
+
+/// Extension methods for collections.
+pub trait CollectionExt<T>
+    where T: IntoIterator
+{
+    /// Consumes the collection and returns its first element.
+    ///
+    /// This method panics if the collection does not have at least one element.
+    fn into_element(self) -> T::Item;
+
+    /// Consumes the collection and returns its last element.
+    ///
+    /// This method panics if the collection does not have at least one element.
+    fn into_last(self) -> T::Item;
+}
+
+impl<T> CollectionExt<T> for T
+    where T: IntoIterator
+{
+    fn into_element(self) -> T::Item {
+        self.into_iter().next().unwrap()
+    }
+
+    fn into_last(self) -> T::Item {
+        self.into_iter().last().unwrap()
+    }
+}

src/ore/lib.rs

@@ -11,12 +11,12 @@
 
 #![forbid(missing_docs)]
 
+pub mod collections;
 pub mod future;
 pub mod log;
 pub mod netio;
 pub mod option;
 pub mod sync;
-pub mod vec;
 
 /// Logs a message to stderr and crashes the process.
 ///

src/ore/vec.rs

@@ -25,7 +25,7 @@ use serde_json::Value as JsonValue;
 use crate::action::{Action, State};
 use crate::parser::BuiltinCommand;
 use ore::future::StreamExt;
-use ore::vec::VecExt;
+use ore::collections::CollectionExt;
 
 pub struct IngestAction {
     topic: String,

src/testdrive/action/kafka.rs

@@ -13,7 +13,7 @@ use std::time::Duration;
 
 use crate::action::{Action, State};
 use crate::parser::{FailSqlCommand, SqlCommand};
-use ore::vec::VecExt;
+use ore::collections::CollectionExt;
 
 pub struct SqlAction {
     cmd: SqlCommand,