Merge pull request #221 from influxdata/crepererum/sandbox-dt-edition

crepererum · web-flow · commit 601396864b7f · 2025-11-27T10:16:51.000Z
feat: check DataType-related conversion
diff --git a/guests/evil/src/complex/many_inputs.rs b/guests/evil/src/complex/many_inputs.rs
@@ -0,0 +1,61 @@
+//! UDF with many inputs.
+
+use std::sync::Arc;
+
+use arrow::datatypes::DataType;
+use datafusion_common::{DataFusionError, Result as DataFusionResult};
+use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
+
+/// UDF with specific signature.
+#[derive(Debug, PartialEq, Eq, Hash)]
+struct SignatureUDF {
+    /// Name.
+    name: &'static str,
+
+    /// The signature.
+    signature: Signature,
+}
+
+impl SignatureUDF {
+    /// Create new UDF.
+    fn new(name: &'static str, signature: Signature) -> Self {
+        Self { name, signature }
+    }
+}
+
+impl ScalarUDFImpl for SignatureUDF {
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        self.name
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> DataFusionResult<DataType> {
+        Ok(DataType::Null)
+    }
+
+    fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
+        Err(DataFusionError::NotImplemented(
+            "invoke_with_args".to_owned(),
+        ))
+    }
+}
+
+/// Returns our evil UDFs.
+///
+/// The passed `source` is ignored.
+#[expect(clippy::unnecessary_wraps, reason = "public API through export! macro")]
+pub(crate) fn udfs(_source: String) -> DataFusionResult<Vec<Arc<dyn ScalarUDFImpl>>> {
+    let limit: usize = std::env::var("limit").unwrap().parse().unwrap();
+
+    Ok(vec![Arc::new(SignatureUDF::new(
+        "input_count",
+        Signature::exact(vec![DataType::Null; limit + 1], Volatility::Immutable),
+    ))])
+}
diff --git a/guests/evil/src/complex/mod.rs b/guests/evil/src/complex/mod.rs
@@ -7,6 +7,9 @@ use datafusion_expr::{
 };
 
 pub(crate) mod error;
+pub(crate) mod many_inputs;
+pub(crate) mod return_type;
+pub(crate) mod return_value;
 pub(crate) mod udf_long_name;
 pub(crate) mod udfs_duplicate_names;
 pub(crate) mod udfs_many;
diff --git a/guests/evil/src/complex/return_type.rs b/guests/evil/src/complex/return_type.rs
@@ -0,0 +1,117 @@
+//! Badness related to return data types.
+
+use std::{collections::HashMap, sync::Arc};
+
+use arrow::datatypes::{DataType, Field};
+use datafusion_common::{DataFusionError, Result as DataFusionResult};
+use datafusion_expr::{
+    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
+};
+
+/// UDF that returns a specific data type.
+#[derive(Debug, PartialEq, Eq, Hash)]
+struct ReturnTypeUDF {
+    /// Name.
+    name: &'static str,
+
+    /// The return type
+    return_type: DataType,
+}
+
+impl ReturnTypeUDF {
+    /// Create new UDF.
+    fn new(name: &'static str, return_type: DataType) -> Self {
+        Self { name, return_type }
+    }
+}
+
+impl ScalarUDFImpl for ReturnTypeUDF {
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        self.name
+    }
+
+    fn signature(&self) -> &Signature {
+        static S: Signature = Signature {
+            type_signature: TypeSignature::Uniform(0, vec![]),
+            volatility: Volatility::Immutable,
+        };
+
+        &S
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> DataFusionResult<DataType> {
+        Ok(self.return_type.clone())
+    }
+
+    fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
+        Err(DataFusionError::NotImplemented(
+            "invoke_with_args".to_owned(),
+        ))
+    }
+}
+
+/// Returns our evil UDFs.
+///
+/// The passed `source` is ignored.
+#[expect(clippy::unnecessary_wraps, reason = "public API through export! macro")]
+pub(crate) fn udfs(_source: String) -> DataFusionResult<Vec<Arc<dyn ScalarUDFImpl>>> {
+    let max_identifier_length: usize = std::env::var("max_identifier_length")
+        .unwrap()
+        .parse()
+        .unwrap();
+    let max_aux_string_length: usize = std::env::var("max_aux_string_length")
+        .unwrap()
+        .parse()
+        .unwrap();
+    let max_depth: usize = std::env::var("max_depth").unwrap().parse().unwrap();
+    let max_complexity: usize = std::env::var("max_complexity").unwrap().parse().unwrap();
+
+    Ok(vec![
+        Arc::new(ReturnTypeUDF::new(
+            "dt_depth",
+            (0..=max_depth).fold(DataType::Int64, |dt, _| {
+                DataType::List(Arc::new(Field::new("f", dt, true)))
+            }),
+        )),
+        Arc::new(ReturnTypeUDF::new(
+            "field_name",
+            DataType::List(Arc::new(Field::new(
+                std::iter::repeat_n('x', max_identifier_length + 1).collect::<String>(),
+                DataType::Int64,
+                true,
+            ))),
+        )),
+        Arc::new(ReturnTypeUDF::new(
+            "field_md_key",
+            DataType::List(Arc::new(
+                Field::new("f", DataType::Int64, true).with_metadata(HashMap::from([(
+                    std::iter::repeat_n('x', max_identifier_length + 1).collect::<String>(),
+                    "value".to_owned(),
+                )])),
+            )),
+        )),
+        Arc::new(ReturnTypeUDF::new(
+            "field_md_value",
+            DataType::List(Arc::new(
+                Field::new("f", DataType::Int64, true).with_metadata(HashMap::from([(
+                    "key".to_owned(),
+                    std::iter::repeat_n('x', max_aux_string_length + 1).collect::<String>(),
+                )])),
+            )),
+        )),
+        Arc::new(ReturnTypeUDF::new(
+            "field_md_items",
+            DataType::List(Arc::new(
+                Field::new("f", DataType::Int64, true).with_metadata(
+                    (0..=max_complexity)
+                        .map(|idx| (idx.to_string(), "value".to_owned()))
+                        .collect(),
+                ),
+            )),
+        )),
+    ])
+}
diff --git a/guests/evil/src/complex/return_value.rs b/guests/evil/src/complex/return_value.rs
@@ -0,0 +1,93 @@
+//! Payloads that return complex values when being invoked.
+
+use std::sync::Arc;
+
+use arrow::{
+    array::{Int64Array, StructArray},
+    datatypes::{DataType, Field},
+};
+use datafusion_common::{DataFusionError, Result as DataFusionResult, ScalarValue};
+use datafusion_expr::{
+    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
+};
+
+use crate::common::DynBox;
+
+/// UDF that returns a specific data type.
+#[derive(Debug, PartialEq, Eq, Hash)]
+struct ReturnValueUDF {
+    /// Name.
+    name: &'static str,
+
+    /// The return value
+    return_value: DynBox<ColumnarValue>,
+}
+
+impl ReturnValueUDF {
+    /// Create new UDF.
+    fn new(name: &'static str, return_value: ColumnarValue) -> Self {
+        Self {
+            name,
+            return_value: DynBox(Box::new(return_value)),
+        }
+    }
+}
+
+impl ScalarUDFImpl for ReturnValueUDF {
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        self.name
+    }
+
+    fn signature(&self) -> &Signature {
+        static S: Signature = Signature {
+            type_signature: TypeSignature::Uniform(0, vec![]),
+            volatility: Volatility::Immutable,
+        };
+
+        &S
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> DataFusionResult<DataType> {
+        Err(DataFusionError::NotImplemented("return_type".to_owned()))
+    }
+
+    fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
+        Ok(self.return_value.clone())
+    }
+}
+
+/// Returns our evil UDFs.
+///
+/// The passed `source` is ignored.
+#[expect(clippy::unnecessary_wraps, reason = "public API through export! macro")]
+pub(crate) fn udfs(_source: String) -> DataFusionResult<Vec<Arc<dyn ScalarUDFImpl>>> {
+    let max_depth: usize = std::env::var("max_depth").unwrap().parse().unwrap();
+
+    Ok(vec![
+        Arc::new(ReturnValueUDF::new(
+            "dt_depth_array",
+            ColumnarValue::Array((0..=max_depth).fold(
+                Arc::new(Int64Array::new(vec![1].into(), None)),
+                |a, _| {
+                    Arc::new(StructArray::from(vec![(
+                        Arc::new(Field::new("a", a.data_type().clone(), true)),
+                        a,
+                    )]))
+                },
+            )),
+        )),
+        Arc::new(ReturnValueUDF::new(
+            "dt_depth_scalar",
+            ColumnarValue::Scalar((0..=max_depth).fold(ScalarValue::Int64(Some(1)), |v, _| {
+                ScalarValue::Struct(Arc::new(StructArray::from(vec![(
+                    Arc::new(Field::new("a", v.data_type(), true)),
+                    v.to_array().unwrap(),
+                )])))
+            })),
+        )),
+    ])
+}
diff --git a/guests/evil/src/lib.rs b/guests/evil/src/lib.rs
@@ -40,6 +40,18 @@ impl Evil {
                 root: Box::new(common::root_empty),
                 udfs: Box::new(complex::error::udfs),
             },
+            "complex::many_inputs" => Self {
+                root: Box::new(common::root_empty),
+                udfs: Box::new(complex::many_inputs::udfs),
+            },
+            "complex::return_type" => Self {
+                root: Box::new(common::root_empty),
+                udfs: Box::new(complex::return_type::udfs),
+            },
+            "complex::return_value" => Self {
+                root: Box::new(common::root_empty),
+                udfs: Box::new(complex::return_value::udfs),
+            },
             "complex::udf_long_name" => Self {
                 root: Box::new(common::root_empty),
                 udfs: Box::new(complex::udf_long_name::udfs),
diff --git a/host/src/conversion/mod.rs b/host/src/conversion/mod.rs
@@ -1,4 +1,6 @@
 //! Conversion routes from/to WIT types.
+use std::collections::HashMap;
+
 use arrow::{
     array::ArrayRef,
     datatypes::{DataType, Field, IntervalUnit, TimeUnit, UnionFields, UnionMode},
@@ -94,7 +96,7 @@ fn check_union_fields(
     let token = token.sub()?;
 
     for (_idx, field) in ufields.iter() {
-        check_field(field, &token)?;
+        check_field(field, &token).context("field")?;
     }
     Ok(())
 }
@@ -167,26 +169,26 @@ fn check_data_type(
         | DataType::FixedSizeList(field, _)
         | DataType::LargeList(field)
         | DataType::LargeListView(field)
-        | DataType::Map(field, _) => check_field(field, &token),
+        | DataType::Map(field, _) => check_field(field, &token).context("field"),
         DataType::Struct(fields) => {
-            for field in fields {
-                check_field(field, &token)?;
+            for (idx, field) in fields.iter().enumerate() {
+                check_field(field, &token).with_context(|| format!("field {idx}"))?;
             }
             Ok(())
         }
         DataType::Union(ufields, umode) => {
-            check_union_fields(ufields, &token)?;
-            check_union_mode(umode, &token)?;
+            check_union_fields(ufields, &token).context("union fields")?;
+            check_union_mode(umode, &token).context("union mode")?;
             Ok(())
         }
         DataType::Dictionary(dt1, dt2) => {
-            check_data_type(dt1, &token)?;
-            check_data_type(dt2, &token)?;
+            check_data_type(dt1, &token).context("key type")?;
+            check_data_type(dt2, &token).context("value type")?;
             Ok(())
         }
         DataType::RunEndEncoded(field1, field2) => {
-            check_field(field1, &token)?;
-            check_field(field2, &token)?;
+            check_field(field1, &token).context("REE run-ends field")?;
+            check_field(field2, &token).context("REE value field")?;
             Ok(())
         }
     }
@@ -196,12 +198,24 @@ fn check_data_type(
 fn check_field(field: &Field, token: &limits::ComplexityToken) -> datafusion_common::Result<()> {
     let token = token.sub()?;
 
-    token.check_identifier(field.name())?;
-    check_data_type(field.data_type(), &token)?;
+    token.check_identifier(field.name()).context("field name")?;
+    check_data_type(field.data_type(), &token).context("field data type")?;
+    check_metadata(field.metadata(), &token).context("field metadata")?;
+
+    Ok(())
+}
+
+/// Check metadata complexity.
+fn check_metadata(
+    md: &HashMap<String, String>,
+    token: &limits::ComplexityToken,
+) -> datafusion_common::Result<()> {
+    let token = token.sub()?;
 
-    for (k, v) in field.metadata() {
-        token.check_identifier(k)?;
-        token.check_aux_string(v)?;
+    for (k, v) in md {
+        let token = token.sub()?;
+        token.check_identifier(k).context("metadata key")?;
+        token.check_aux_string(v).context("metadata value")?;
     }
 
     Ok(())
diff --git a/host/src/lib.rs b/host/src/lib.rs
@@ -620,7 +620,8 @@ impl WasmScalarUdf {
                     Some(&store_guard.data().stderr.contents()),
                 )?;
             ComplexityToken::new(permissions.trusted_data_limits.clone())?
-                .check_identifier(&name)?;
+                .check_identifier(&name)
+                .context("UDF name")?;
             if !names_seen.insert(name.clone()) {
                 return Err(DataFusionError::External(
                     format!("non-unique UDF name: '{name}'").into(),
diff --git a/host/tests/integration_tests/evil/complex.rs b/host/tests/integration_tests/evil/complex.rs
