mtrlz/sql: support coalesce function

Author: benesch

src/materialize/dataflow/func.rs

@@ -395,6 +395,26 @@ impl UnaryFunc {
     }
 }
 
+pub fn coalesce(datums: Vec<Datum>) -> Datum {
+    datums
+        .into_iter()
+        .find(|d| !d.is_null())
+        .unwrap_or(Datum::Null)
+}
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize, Hash)]
+pub enum VariadicFunc {
+    Coalesce,
+}
+
+impl VariadicFunc {
+    pub fn func(self) -> fn(Vec<Datum>) -> Datum {
+        match self {
+            VariadicFunc::Coalesce => coalesce,
+        }
+    }
+}
+
 // TODO(jamii) be careful about overflow in sum/avg
 // see https://timely.zulipchat.com/#narrow/stream/186635-engineering/topic/additional.20work/near/163507435
 

src/materialize/dataflow/render.rs

@@ -466,6 +466,10 @@ fn eval_expr(expr: &Expr, datum: &Datum) -> Datum {
             let datum2 = eval_expr(expr2, datum);
             (func.func())(datum1, datum2)
         }
+        Expr::CallVariadic { func, exprs } => {
+            let datums = exprs.iter().map(|e| eval_expr(e, datum)).collect();
+            (func.func())(datums)
+        }
         Expr::If { cond, then, els } => match eval_expr(cond, datum) {
             Datum::True => eval_expr(then, datum),
             Datum::False | Datum::Null => eval_expr(els, datum),

src/materialize/dataflow/types.rs

@@ -6,7 +6,7 @@
 use serde::{Deserialize, Serialize};
 use url::Url;
 
-use super::func::{AggregateFunc, BinaryFunc, UnaryFunc};
+use super::func::{AggregateFunc, BinaryFunc, UnaryFunc, VariadicFunc};
 use crate::repr::{Datum, Type};
 
 /// System-wide update type.
@@ -231,11 +231,11 @@ pub enum Expr {
         then: Box<Expr>,
         els: Box<Expr>,
     },
-    // /// A function call that takes an arbitrary number of arguments.
-    // CallMany {
-    //     fn: fn(Vec<Datum>) -> Datum,
-    //     exprs: Vec<Expr>,
-    // }
+    /// A function call that takes an arbitrary number of arguments.
+    CallVariadic {
+        func: VariadicFunc,
+        exprs: Vec<Expr>,
+    },
 }
 
 impl Expr {

src/materialize/sql/mod.rs

@@ -22,7 +22,7 @@ use std::iter::FromIterator;
 use std::net::{SocketAddr, ToSocketAddrs};
 use url::Url;
 
-use crate::dataflow::func::{AggregateFunc, BinaryFunc, UnaryFunc};
+use crate::dataflow::func::{AggregateFunc, BinaryFunc, UnaryFunc, VariadicFunc};
 use crate::dataflow::{
     Aggregate, Dataflow, Expr, KafkaSinkConnector, KafkaSourceConnector, LocalSourceConnector,
     Plan, Sink, SinkConnector, Source, SourceConnector, View,
@@ -1100,7 +1100,7 @@ impl Planner {
             }
             let (i, typ) = plan.resolve_func(name);
             let expr = Expr::Column(i, Box::new(Expr::Ambient));
-            return Ok((expr, typ.clone()))
+            return Ok((expr, typ.clone()));
         }
 
         match ident.as_str() {
@@ -1116,11 +1116,45 @@ impl Planner {
                     FType::Float64 => UnaryFunc::AbsFloat64,
                     _ => bail!("abs does not accept arguments of type {:?}", typ),
                 };
-                Ok((Expr::CallUnary {
+                let expr = Expr::CallUnary {
                     func,
                     expr: Box::new(expr),
-                }, typ))
+                };
+                Ok((expr, typ))
             }
+
+            "coalesce" => {
+                if args.is_empty() {
+                    bail!("coalesce requires at least one argument");
+                }
+                let mut exprs = Vec::new();
+                let mut result_type: Option<Type> = None;
+                for arg in args {
+                    let (expr, typ) = self.plan_expr(ctx, arg, plan)?;
+                    match &result_type {
+                        Some(result_type) => {
+                            if result_type.ftype != typ.ftype {
+                                bail!(
+                                    "COALESCE does not have uniform argument type: {:?} vs {:?}",
+                                    result_type.ftype,
+                                    typ.ftype
+                                )
+                            }
+                        }
+                        None => result_type = Some(typ),
+                    }
+                    exprs.push(expr);
+                }
+                // args is known to be non-empty, and therefore result_type must
+                // be non-None after the loop above.
+                let result_type = result_type.unwrap();
+                let expr = Expr::CallVariadic {
+                    func: VariadicFunc::Coalesce,
+                    exprs,
+                };
+                Ok((expr, result_type))
+            }
+
             _ => bail!("unsupported function: {}", ident),
         }
     }

src/ore/collections.rs

@@ -7,7 +7,8 @@
 
 /// Extension methods for collections.
 pub trait CollectionExt<T>
-    where T: IntoIterator
+where
+    T: IntoIterator,
 {
     /// Consumes the collection and returns its first element.
     ///
@@ -21,7 +22,8 @@ pub trait CollectionExt<T>
 }
 
 impl<T> CollectionExt<T> for T
-    where T: IntoIterator
+where
+    T: IntoIterator,
 {
     fn into_element(self) -> T::Item {
         self.into_iter().next().unwrap()

src/testdrive/action/kafka.rs

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