chore(cubestore): Upgrade DF: remove coalesce UDF, make coalesce test fixes, handle DataType::Null in batches_to_dataframe

Author: srh

rust/cubestore/cubestore-sql-tests/src/tests.rs

@@ -1729,12 +1729,11 @@ async fn coalesce(service: Box<dyn SqlClient>) {
         .await
         .unwrap();
     assert_eq!(to_rows(&r), vec![vec![TableValue::Int(1)]]);
-    // TODO: the type should be 'int' here. Hopefully not a problem in practice.
     let r = service
         .exec_query("SELECT coalesce(NULL, 2, 3)")
         .await
         .unwrap();
-    assert_eq!(to_rows(&r), vec![vec![TableValue::String("2".to_string())]]);
+    assert_eq!(to_rows(&r), vec![vec![TableValue::Int(2)]]);
     let r = service
         .exec_query("SELECT coalesce(NULL, NULL, NULL)")
         .await
@@ -1753,20 +1752,11 @@ async fn coalesce(service: Box<dyn SqlClient>) {
             vec![TableValue::Null],
         ]
     );
-    // Coerces all args to text.
-    let r = service
+    // Type mismatch
+    service
         .exec_query("SELECT coalesce(n, v, s) FROM s.Data ORDER BY 1")
         .await
-        .unwrap();
-    assert_eq!(
-        to_rows(&r),
-        vec![
-            vec![TableValue::String("1".to_string())],
-            vec![TableValue::String("3".to_string())],
-            vec![TableValue::String("baz".to_string())],
-            vec![TableValue::Null],
-        ]
-    );
+        .unwrap_err();
 
     let r = service
         .exec_query("SELECT coalesce(n+1,v+1,0) FROM s.Data ORDER BY 1")

rust/cubestore/cubestore/src/queryplanner/coalesce.rs

@@ -15,7 +15,6 @@ mod topk;
 pub mod trace_data_loaded;
 pub use topk::MIN_TOPK_STREAM_ROWS;
 use udfs::{aggregate_udf_by_kind, registerable_aggregate_udfs, registerable_scalar_udfs};
-mod coalesce;
 mod filter_by_key_range;
 mod flatten_union;
 pub mod info_schema;
@@ -488,7 +487,6 @@ impl ContextProvider for MetaStoreSchemaProvider {
         // TODO upgrade DF
         let kind = match name {
             "cardinality" | "CARDINALITY" => CubeScalarUDFKind::HllCardinality,
-            // "coalesce" | "COALESCE" => CubeScalarUDFKind::Coalesce,
             "unix_timestamp" | "UNIX_TIMESTAMP" => CubeScalarUDFKind::UnixTimestamp,
             "date_add" | "DATE_ADD" => CubeScalarUDFKind::DateAdd,
             "date_sub" | "DATE_SUB" => CubeScalarUDFKind::DateSub,

rust/cubestore/cubestore/src/queryplanner/mod.rs

@@ -24,8 +24,8 @@ use async_trait::async_trait;
 use core::fmt;
 use datafusion::arrow::array::{
     make_array, Array, ArrayRef, BinaryArray, BooleanArray, Decimal128Array, Float64Array,
-    Int16Array, Int32Array, Int64Array, MutableArrayData, StringArray, TimestampMicrosecondArray,
-    TimestampNanosecondArray, UInt16Array, UInt32Array, UInt64Array,
+    Int16Array, Int32Array, Int64Array, MutableArrayData, NullArray, StringArray,
+    TimestampMicrosecondArray, TimestampNanosecondArray, UInt16Array, UInt32Array, UInt64Array,
 };
 use datafusion::arrow::compute::SortOptions;
 use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
@@ -92,7 +92,10 @@ use std::sync::Arc;
 use std::time::SystemTime;
 use tracing::{instrument, Instrument};
 
-use super::udfs::{aggregate_udf_by_kind, registerable_aggregate_udfs, registerable_arc_aggregate_udfs, registerable_arc_scalar_udfs, CubeAggregateUDFKind};
+use super::udfs::{
+    aggregate_udf_by_kind, registerable_aggregate_udfs, registerable_arc_aggregate_udfs,
+    registerable_arc_scalar_udfs, CubeAggregateUDFKind,
+};
 
 #[automock]
 #[async_trait]
@@ -1926,6 +1929,13 @@ pub fn batches_to_dataframe(batches: Vec<RecordBatch>) -> Result<DataFrame, Cube
                         });
                     }
                 }
+                DataType::Null => {
+                    // Force the cast, just because.
+                    let _ = array.as_any().downcast_ref::<NullArray>().unwrap();
+                    for i in 0..num_rows {
+                        rows[i].push(TableValue::Null);
+                    }
+                }
                 x => panic!("Unsupported data type: {:?}", x),
             }
         }
@@ -1962,6 +1972,8 @@ pub fn arrow_to_column_type(arrow_type: DataType) -> Result<ColumnType, CubeErro
         | DataType::UInt16
         | DataType::UInt32
         | DataType::UInt64 => Ok(ColumnType::Int),
+        // This fn is only used for converting to DataFrame, and cubesql does this (as if that's a reason)
+        DataType::Null => Ok(ColumnType::String),
         x => Err(CubeError::internal(format!("unsupported type {:?}", x))),
     }
 }

rust/cubestore/cubestore/src/queryplanner/query_executor.rs

@@ -1,4 +1,3 @@
-use crate::queryplanner::coalesce::SUPPORTED_COALESCE_TYPES;
 use crate::queryplanner::hll::{Hll, HllUnion};
 use crate::CubeError;
 use chrono::{Datelike, Duration, Months, NaiveDateTime, TimeZone, Utc};
@@ -26,7 +25,6 @@ use std::sync::Arc;
 #[derive(Copy, Clone, Debug, Serialize, Deserialize)]
 pub enum CubeScalarUDFKind {
     HllCardinality, // cardinality(), accepting the HyperLogLog sketches.
-    // Coalesce,
     UnixTimestamp,
     DateAdd,
     DateSub,
@@ -36,7 +34,6 @@ pub enum CubeScalarUDFKind {
 pub fn scalar_udf_by_kind(k: CubeScalarUDFKind) -> Arc<ScalarUDF> {
     match k {
         CubeScalarUDFKind::HllCardinality => Arc::new(HllCardinality::descriptor()),
-        // CubeScalarUDFKind::Coalesce => Box::new(Coalesce {}),
         CubeScalarUDFKind::UnixTimestamp => {
             Arc::new(ScalarUDF::new_from_impl(UnixTimestamp::new()))
         }
@@ -67,9 +64,6 @@ pub fn scalar_kind_by_name(n: &str) -> Option<CubeScalarUDFKind> {
     if n == "CARDINALITY" {
         return Some(CubeScalarUDFKind::HllCardinality);
     }
-    // if n == "COALESCE" {
-    //     return Some(CubeScalarUDFKind::Coalesce);
-    // }
     if n == "UNIX_TIMESTAMP" {
         return Some(CubeScalarUDFKind::UnixTimestamp);
     }
@@ -83,7 +77,7 @@ pub fn scalar_kind_by_name(n: &str) -> Option<CubeScalarUDFKind> {
         return Some(CubeScalarUDFKind::DateBin);
     }
     // TODO upgrade DF: Remove this (once we are no longer in flux about naming casing of UDFs and UDAFs).
-    if ["CARDINALITY", /* "COALESCE", */ "UNIX_TIMESTAMP", "DATE_ADD", "DATE_SUB", "DATE_BIN"].contains(&(&n.to_ascii_uppercase() as &str)) {
+    if ["CARDINALITY", "UNIX_TIMESTAMP", "DATE_ADD", "DATE_SUB", "DATE_BIN"].contains(&(&n.to_ascii_uppercase() as &str)) {
         panic!("scalar_kind_by_name failing on '{}' due to uppercase/lowercase mixup", n);
     }
     return None;
@@ -126,40 +120,6 @@ pub fn aggregate_kind_by_name(n: &str) -> Option<CubeAggregateUDFKind> {
 // The rest of the file are implementations of the various functions that we have.
 // TODO: add custom type and use it instead of `Binary` for HLL columns.
 
-// TODO upgrade DF - remove?
-// struct Coalesce {}
-// impl Coalesce {
-//     fn signature() -> Signature {
-//         Signature::Variadic(SUPPORTED_COALESCE_TYPES.to_vec())
-//     }
-// }
-// impl CubeScalarUDF for Coalesce {
-//     fn kind(&self) -> CubeScalarUDFKind {
-//         CubeScalarUDFKind::Coalesce
-//     }
-//
-//     fn name(&self) -> &str {
-//         "COALESCE"
-//     }
-//
-//     fn descriptor(&self) -> ScalarUDF {
-//         return ScalarUDF {
-//             name: self.name().to_string(),
-//             signature: Self::signature(),
-//             return_type: Arc::new(|inputs| {
-//                 if inputs.is_empty() {
-//                     return Err(DataFusionError::Plan(
-//                         "COALESCE requires at least 1 argument".to_string(),
-//                     ));
-//                 }
-//                 let ts = type_coercion::data_types(inputs, &Self::signature())?;
-//                 Ok(Arc::new(ts[0].clone()))
-//             }),
-//             fun: Arc::new(coalesce),
-//         };
-//     }
-// }
-
 #[derive(Debug)]
 struct UnixTimestamp {
     signature: Signature,