paleolimbot commented on code in PR #24939:
URL: https://github.com/apache/datafusion/pull/24939#discussion_r3930558457
##########
datafusion/core/tests/user_defined/user_defined_scalar_functions.rs:
##########
@@ -2255,3 +2255,99 @@ async fn test_extension_metadata_preserve_in_subquery()
-> Result<()> {
assert!(!df.collect().await?.is_empty());
Ok(())
}
+
+/// Uncorrelated scalar subqueries must keep Arrow field metadata on the
+/// physical expression so UDFs that distinguish extension types still
type-check.
+/// https://github.com/apache/datafusion/issues/24933
+#[tokio::test]
+async fn test_extension_metadata_preserve_in_uncorrelated_scalar_subquery() ->
Result<()>
+{
+ #[derive(Debug, PartialEq, Eq, Hash)]
+ struct MetadataRequired {
+ signature: Signature,
+ }
+
+ impl Default for MetadataRequired {
+ fn default() -> Self {
+ Self {
+ signature: Signature::user_defined(Volatility::Immutable),
+ }
+ }
+ }
+
+ impl ScalarUDFImpl for MetadataRequired {
+ fn name(&self) -> &str {
+ "metadata_required"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn coerce_types(&self, arg_types: &[DataType]) ->
Result<Vec<DataType>> {
+ Ok(arg_types.to_vec())
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+ unreachable!("return_field_from_args is implemented")
+ }
+
+ fn return_field_from_args(&self, args: ReturnFieldArgs) ->
Result<FieldRef> {
+ for (i, field) in args.arg_fields.iter().enumerate() {
+ if !field.metadata().contains_key("ARROW:extension:name") {
+ return exec_err!(
+ "argument {i} lost ARROW:extension:name;
field={field:?} metadata={:?}",
+ field.metadata()
+ );
+ }
+ }
+
+ Ok(Field::new("metadata_required", DataType::Boolean, true).into())
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) ->
Result<ColumnarValue> {
+ Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(
+ args.arg_fields
+ .iter()
+ .all(|field|
field.metadata().contains_key("ARROW:extension:name")),
+ ))))
+ }
+ }
+
+ let schema = Schema::new(vec![
+ Field::new("id", DataType::Int64, false),
+ Field::new("geometry", DataType::Utf8,
false).with_metadata(HashMap::from([(
+ "ARROW:extension:name".to_string(),
+ "example.extension".to_string(),
+ )])),
+ ]);
+
+ let batch = RecordBatch::try_new(
+ schema.clone().into(),
+ vec![
+ create_array!(Int64, [1, 2]),
+ create_array!(Utf8, [Some("a"), Some("b")]),
+ ],
+ )?;
+
+ let ctx = SessionContext::new();
+ ctx.register_batch("l", batch.clone())?;
+ ctx.register_batch("r", batch)?;
+ ctx.register_udf(MetadataRequired::default().into());
+
+ let df = ctx
+ .sql(
+ "
+ SELECT id
+ FROM l
+ WHERE metadata_required(
+ l.geometry,
+ (SELECT r.geometry FROM r WHERE r.id = 1)
+ )
Review Comment:
Is there a good way to get a query that exercises this into a `.slt` file as
well? Maybe `datafusion/sqllogictest/test_files/metadata.slt`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
