Merge pull request #53 from influxdata/crepererum/colvar_helpers

test: avoid boilerplate when working with `ColumnarValue`

Author: crepererum

host/tests/integration_tests/mod.rs

@@ -1,2 +1,3 @@
 mod python;
 mod rust;
+mod test_utils;

host/tests/integration_tests/python.rs

@@ -2,9 +2,11 @@ use std::sync::Arc;
 
 use arrow::datatypes::{DataType, Field};
 use datafusion_common::{ScalarValue, assert_contains};
-use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
+use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
 use datafusion_udf_wasm_host::{WasmComponentPrecompiled, WasmScalarUdf};
 
+use crate::integration_tests::test_utils::ColumnarValueExt;
+
 #[tokio::test(flavor = "multi_thread")]
 async fn test() {
     let data = tokio::fs::read(format!(
@@ -28,17 +30,15 @@ async fn test() {
 
     assert_eq!(udf.return_type(&[]).unwrap(), DataType::Utf8,);
 
-    let ColumnarValue::Scalar(scalar) = udf
+    let scalar = udf
         .invoke_with_args(ScalarFunctionArgs {
             args: vec![],
             arg_fields: vec![],
             number_rows: 3,
             return_field: Arc::new(Field::new("r", DataType::Utf8, true)),
         })
         .unwrap()
-    else {
-        panic!("should be a scalar");
-    };
+        .unwrap_scalar();
     let ScalarValue::Utf8(s) = scalar else {
         panic!("scalar should be UTF8");
     };

host/tests/integration_tests/rust.rs

@@ -8,6 +8,8 @@ use datafusion_common::ScalarValue;
 use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
 use datafusion_udf_wasm_host::{WasmComponentPrecompiled, WasmScalarUdf};
 
+use crate::integration_tests::test_utils::ColumnarValueExt;
+
 #[tokio::test(flavor = "multi_thread")]
 async fn test_add_one() {
     let data = tokio::fs::read(format!(
@@ -38,7 +40,7 @@ async fn test_add_one() {
         @"Error during planning: add_one expects exactly one argument",
     );
 
-    let ColumnarValue::Array(array) = udf
+    let array = udf
         .invoke_with_args(ScalarFunctionArgs {
             args: vec![ColumnarValue::Array(Arc::new(Int32Array::from_iter([
                 Some(3),
@@ -50,23 +52,19 @@ async fn test_add_one() {
             return_field: Arc::new(Field::new("r", DataType::Int32, true)),
         })
         .unwrap()
-    else {
-        panic!("should be an array")
-    };
+        .unwrap_array();
     assert_eq!(
         array.as_ref(),
         &Int32Array::from_iter([Some(4), None, Some(2)]) as &dyn Array,
     );
-    let ColumnarValue::Scalar(scalar) = udf
+    let scalar = udf
         .invoke_with_args(ScalarFunctionArgs {
             args: vec![ColumnarValue::Scalar(ScalarValue::Int32(Some(3)))],
             arg_fields: vec![Arc::new(Field::new("a1", DataType::Int32, true))],
             number_rows: 3,
             return_field: Arc::new(Field::new("r", DataType::Int32, true)),
         })
         .unwrap()
-    else {
-        panic!("should be a scalar")
-    };
+        .unwrap_scalar();
     assert_eq!(scalar, ScalarValue::Int32(Some(4)));
 }

host/tests/integration_tests/test_utils.rs

@@ -0,0 +1,38 @@
+use arrow::array::ArrayRef;
+use datafusion_common::ScalarValue;
+use datafusion_expr::ColumnarValue;
+
+/// Extension trait for [`ColumnarValue`] for easier testing.
+pub(crate) trait ColumnarValueExt {
+    /// Extracts [`ColumnarValue::Array`] variant.
+    ///
+    /// # Panic
+    /// Panics if this is not an array.
+    #[track_caller]
+    fn unwrap_array(self) -> ArrayRef;
+
+    /// Extracts [`ColumnarValue::Scalar`] variant.
+    ///
+    /// # Panic
+    /// Panics if this is not an scalar.
+    #[track_caller]
+    fn unwrap_scalar(self) -> ScalarValue;
+}
+
+impl ColumnarValueExt for ColumnarValue {
+    #[track_caller]
+    fn unwrap_array(self) -> ArrayRef {
+        match self {
+            Self::Array(array) => array,
+            Self::Scalar(_) => panic!("expected an array but got a scalar"),
+        }
+    }
+
+    #[track_caller]
+    fn unwrap_scalar(self) -> ScalarValue {
+        match self {
+            Self::Array(_) => panic!("expected a scalar but got an array"),
+            Self::Scalar(scalar_value) => scalar_value,
+        }
+    }
+}