Support more cast_to_variant arguments

Author: friendlymatthew

src/cast_to_variant.rs

@@ -13,7 +13,7 @@ use datafusion::{
 use parquet_variant::Variant;
 use parquet_variant_compute::{VariantArray, VariantArrayBuilder};
 
-use crate::shared::{try_field_as_binary, try_parse_binary_columnar, try_parse_binary_scalar};
+use crate::shared::{try_parse_binary_columnar, try_parse_binary_scalar};
 
 #[derive(Debug, Hash, PartialEq, Eq)]
 pub struct CastToVariantUdf {
@@ -23,48 +23,26 @@ pub struct CastToVariantUdf {
 impl Default for CastToVariantUdf {
     fn default() -> Self {
         Self {
-            signature: Signature::new(TypeSignature::Any(2), Volatility::Immutable),
+            signature: Signature::new(TypeSignature::VariadicAny, Volatility::Immutable),
         }
     }
 }
 
-impl ScalarUDFImpl for CastToVariantUdf {
-    fn as_any(&self) -> &dyn std::any::Any {
-        self
-    }
-
-    fn name(&self) -> &str {
-        "cast_to_variant"
-    }
-
-    fn signature(&self) -> &Signature {
-        &self.signature
-    }
-
-    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
-        Ok(DataType::Struct(Fields::from(vec![
-            Field::new("metadata", DataType::BinaryView, false),
-            Field::new("value", DataType::BinaryView, true),
-        ])))
-    }
+fn build_variant_array<'m, 'v, T: Into<Variant<'m, 'v>>>(
+    value_opt: Option<T>,
+) -> Result<ColumnarValue> {
+    let variant_array = VariantArray::from_iter([value_opt.map(|v| v.into())]).into();
 
-    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
-        match args.arg_fields.as_slice() {
-            [metadata_field, variant_field] => {
-                try_field_as_binary(metadata_field.as_ref())?;
-                try_field_as_binary(variant_field.as_ref())?;
-            }
-            _ => {
-                // right now, let's only support (BinaryViewArray, BinaryViewArray)
-                // but I don't see why we couldn't call cast_to_variant(string_column) -> VariantArray...
-                return exec_err!("unsupported, expected 2 arguments");
-            }
-        }
-
-        let [metadata_argument, variant_argument] = args.args.as_slice() else {
-            return exec_err!("expected 2 arguments");
-        };
+    Ok(ColumnarValue::Scalar(ScalarValue::Struct(Arc::new(
+        variant_array,
+    ))))
+}
 
+impl CastToVariantUdf {
+    fn from_metadata_value(
+        metadata_argument: &ColumnarValue,
+        variant_argument: &ColumnarValue,
+    ) -> Result<ColumnarValue> {
         let out = match (metadata_argument, variant_argument) {
             (ColumnarValue::Array(metadata_array), ColumnarValue::Array(value_array)) => {
                 if metadata_array.len() != value_array.len() {
@@ -154,6 +132,73 @@ impl ScalarUDFImpl for CastToVariantUdf {
 
         Ok(out)
     }
+
+    fn from_array(_array: &ArrayRef) -> Result<ColumnarValue> {
+        todo!()
+    }
+
+    fn from_scalar_value(scalar_value: &ScalarValue) -> Result<ColumnarValue> {
+        match scalar_value {
+            ScalarValue::Null => build_variant_array(Some(Variant::Null)),
+            // String values
+            ScalarValue::Utf8(string_opt)
+            | ScalarValue::Utf8View(string_opt)
+            | ScalarValue::LargeUtf8(string_opt) => {
+                build_variant_array(string_opt.as_ref().map(|s| s.as_str()))
+            }
+            // Binary values
+            ScalarValue::Binary(binary_opt)
+            | ScalarValue::BinaryView(binary_opt)
+            | ScalarValue::LargeBinary(binary_opt) => {
+                build_variant_array(binary_opt.as_ref().map(|b| b.as_slice()))
+            }
+            // Boolean
+            ScalarValue::Boolean(b) => build_variant_array(b.as_ref().map(|b| *b)),
+            // Numbers
+            ScalarValue::Int8(i) => build_variant_array(i.as_ref().map(|i| *i)),
+            ScalarValue::Int16(i) => build_variant_array(i.as_ref().map(|i| *i)),
+            ScalarValue::Int32(i) => build_variant_array(i.as_ref().map(|i| *i)),
+            ScalarValue::Int64(i) => build_variant_array(i.as_ref().map(|i| *i)),
+            ScalarValue::UInt8(i) => build_variant_array(i.as_ref().map(|i| *i)),
+            ScalarValue::UInt16(i) => build_variant_array(i.as_ref().map(|i| *i)),
+            ScalarValue::UInt32(i) => build_variant_array(i.as_ref().map(|i| *i)),
+            ScalarValue::UInt64(i) => build_variant_array(i.as_ref().map(|i| *i)),
+
+            _ => todo!(),
+        }
+    }
+}
+
+impl ScalarUDFImpl for CastToVariantUdf {
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        "cast_to_variant"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+        Ok(DataType::Struct(Fields::from(vec![
+            Field::new("metadata", DataType::BinaryView, false),
+            Field::new("value", DataType::BinaryView, true),
+        ])))
+    }
+
+    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        match args.args.as_slice() {
+            [metadata_value, variant_value] => {
+                Self::from_metadata_value(metadata_value, variant_value)
+            }
+            [ColumnarValue::Scalar(scalar_value)] => Self::from_scalar_value(scalar_value),
+            [ColumnarValue::Array(array)] => Self::from_array(array),
+            _ => exec_err!("unrecognized argument"),
+        }
+    }
 }
 
 #[cfg(test)]

src/shared.rs

@@ -26,7 +26,7 @@ pub fn try_field_as_variant_array(field: &Field) -> Result<()> {
     Ok(())
 }
 
-pub fn try_field_as_binary(field: &Field) -> Result<()> {
+pub fn _try_field_as_binary(field: &Field) -> Result<()> {
     match field.data_type() {
         DataType::Binary | DataType::BinaryView | DataType::LargeBinary => {}
         unsupported => return exec_err!("expected binary field, got {unsupported} field"),