M-Tesla commented on code in PR #24939:
URL: https://github.com/apache/datafusion/pull/24939#discussion_r3933647368
##########
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:
Thank you for the review, and for pointing at `metadata.slt` — that was the
right place.
I added an uncorrelated scalar subquery case there using `arrow_metadata`,
next to the existing CAST/window metadata checks. Locally `cargo test --test
sqllogictests -- metadata.slt` passes with the fix; if the subquery drops field
metadata, the map query returns `{}` instead of `{metadata_key: the id field}`.
--
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]
