Skip to content

Commit c74dd77

Browse files
paleolimbotpaleolimbot
committed
try the generated answer
1 parent a0e9887 commit c74dd77

4 files changed

Lines changed: 260 additions & 22 deletions

File tree

‎Cargo.lock‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎datafusion/expr/src/expr_schema.rs‎

Lines changed: 152 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -69,18 +69,43 @@ pub trait ExprSchemable {
6969
-> Result<(DataType, bool)>;
7070
}
7171

72-
/// Derives the output field for a cast expression from the source field.
72+
/// Derives the output field for a cast expression from the source and target fields.
73+
///
74+
/// Metadata handling:
75+
/// - Source field metadata is propagated by default
76+
/// - Extension type metadata keys (`ARROW:extension:name` and `ARROW:extension:metadata`)
77+
/// are taken from the target field, overwriting any extension metadata from the source.
78+
/// This ensures casting does not incorrectly propagate extension type identity.
79+
///
7380
/// For `TryCast`, `force_nullable` is `true` since a failed cast returns NULL.
7481
fn cast_output_field(
7582
source_field: &FieldRef,
76-
target_type: &DataType,
83+
target_field: &FieldRef,
7784
force_nullable: bool,
7885
) -> Arc<Field> {
86+
use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY};
87+
88+
// Start with source metadata
89+
let mut metadata = source_field.metadata().clone();
90+
91+
// Remove any extension type metadata from source - these should not propagate through casts
92+
metadata.remove(EXTENSION_TYPE_NAME_KEY);
93+
metadata.remove(EXTENSION_TYPE_METADATA_KEY);
94+
95+
// Add extension type metadata from the target field if present
96+
let target_metadata = target_field.metadata();
97+
if let Some(name) = target_metadata.get(EXTENSION_TYPE_NAME_KEY) {
98+
metadata.insert(EXTENSION_TYPE_NAME_KEY.to_string(), name.clone());
99+
}
100+
if let Some(ext_meta) = target_metadata.get(EXTENSION_TYPE_METADATA_KEY) {
101+
metadata.insert(EXTENSION_TYPE_METADATA_KEY.to_string(), ext_meta.clone());
102+
}
103+
79104
let mut f = source_field
80105
.as_ref()
81106
.clone()
82-
.with_data_type(target_type.clone())
83-
.with_metadata(source_field.metadata().clone());
107+
.with_data_type(target_field.data_type().clone())
108+
.with_metadata(metadata);
84109
if force_nullable {
85110
f = f.with_nullable(true);
86111
}
@@ -595,20 +620,16 @@ impl ExprSchemable for Expr {
595620
func.return_field_from_args(args)
596621
}
597622
// _ => Ok((self.get_type(schema)?, self.nullable(schema)?)),
598-
Expr::Cast(Cast { expr, field }) => {
599-
expr.to_field(schema).map(|(_table_ref, src)| {
600-
cast_output_field(&src, field.data_type(), false)
601-
})
602-
}
623+
Expr::Cast(Cast { expr, field }) => expr
624+
.to_field(schema)
625+
.map(|(_table_ref, src)| cast_output_field(&src, field, false)),
603626
Expr::Placeholder(Placeholder {
604627
id: _,
605628
field: Some(field),
606629
}) => Ok(Arc::clone(field).renamed(&schema_name)),
607-
Expr::TryCast(TryCast { expr, field }) => {
608-
expr.to_field(schema).map(|(_table_ref, src)| {
609-
cast_output_field(&src, field.data_type(), true)
610-
})
611-
}
630+
Expr::TryCast(TryCast { expr, field }) => expr
631+
.to_field(schema)
632+
.map(|(_table_ref, src)| cast_output_field(&src, field, true)),
612633
Expr::LambdaVariable(LambdaVariable {
613634
field: Some(field), ..
614635
}) => Ok(Arc::clone(field).renamed(&schema_name)),
@@ -1208,4 +1229,121 @@ mod tests {
12081229

12091230
assert_eq!(meta, expr.metadata(&schema).unwrap());
12101231
}
1232+
1233+
#[test]
1234+
fn test_cast_extension_type_metadata() {
1235+
use crate::expr::Cast;
1236+
use arrow_schema::extension::{
1237+
EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY,
1238+
};
1239+
1240+
// Create a schema with a field that has extension type metadata
1241+
let mut source_meta = HashMap::new();
1242+
source_meta.insert(
1243+
EXTENSION_TYPE_NAME_KEY.to_string(),
1244+
"arrow.uuid".to_string(),
1245+
);
1246+
source_meta.insert("custom_key".to_string(), "custom_value".to_string());
1247+
1248+
let source_field = Field::new("foo", DataType::FixedSizeBinary(16), false)
1249+
.with_metadata(source_meta);
1250+
1251+
let schema = MockExprSchema::new()
1252+
.with_data_type(DataType::FixedSizeBinary(16))
1253+
.with_metadata(FieldMetadata::from(source_field.metadata().clone()));
1254+
1255+
// Test 1: Cast to a type without extension metadata strips extension metadata
1256+
// but preserves non-extension metadata
1257+
let cast_expr = Expr::Cast(Cast {
1258+
expr: Box::new(col("foo")),
1259+
field: Arc::new(Field::new("", DataType::Utf8, true)),
1260+
});
1261+
1262+
let (_, result_field) = cast_expr.to_field(&schema).unwrap();
1263+
assert!(
1264+
result_field
1265+
.metadata()
1266+
.get(EXTENSION_TYPE_NAME_KEY)
1267+
.is_none(),
1268+
"Extension type name should be stripped when target has no extension metadata"
1269+
);
1270+
assert_eq!(
1271+
result_field.metadata().get("custom_key"),
1272+
Some(&"custom_value".to_string()),
1273+
"Non-extension metadata should be preserved"
1274+
);
1275+
1276+
// Test 2: Cast to a field with different extension type replaces extension metadata
1277+
let mut target_meta = HashMap::new();
1278+
target_meta.insert(
1279+
EXTENSION_TYPE_NAME_KEY.to_string(),
1280+
"arrow.json".to_string(),
1281+
);
1282+
target_meta.insert(EXTENSION_TYPE_METADATA_KEY.to_string(), "{}".to_string());
1283+
1284+
let target_field =
1285+
Field::new("", DataType::Utf8, true).with_metadata(target_meta);
1286+
1287+
let cast_expr = Expr::Cast(Cast {
1288+
expr: Box::new(col("foo")),
1289+
field: Arc::new(target_field),
1290+
});
1291+
1292+
let (_, result_field) = cast_expr.to_field(&schema).unwrap();
1293+
assert_eq!(
1294+
result_field.metadata().get(EXTENSION_TYPE_NAME_KEY),
1295+
Some(&"arrow.json".to_string()),
1296+
"Extension type name should come from target field"
1297+
);
1298+
assert_eq!(
1299+
result_field.metadata().get(EXTENSION_TYPE_METADATA_KEY),
1300+
Some(&"{}".to_string()),
1301+
"Extension type metadata should come from target field"
1302+
);
1303+
assert_eq!(
1304+
result_field.metadata().get("custom_key"),
1305+
Some(&"custom_value".to_string()),
1306+
"Non-extension metadata should still be preserved from source"
1307+
);
1308+
}
1309+
1310+
#[test]
1311+
fn test_try_cast_extension_type_metadata() {
1312+
use crate::expr::TryCast;
1313+
use arrow_schema::extension::EXTENSION_TYPE_NAME_KEY;
1314+
1315+
// Create a schema with a field that has extension type metadata
1316+
let mut source_meta = HashMap::new();
1317+
source_meta.insert(
1318+
EXTENSION_TYPE_NAME_KEY.to_string(),
1319+
"arrow.uuid".to_string(),
1320+
);
1321+
1322+
let source_field = Field::new("foo", DataType::FixedSizeBinary(16), false)
1323+
.with_metadata(source_meta);
1324+
1325+
let schema = MockExprSchema::new()
1326+
.with_data_type(DataType::FixedSizeBinary(16))
1327+
.with_metadata(FieldMetadata::from(source_field.metadata().clone()));
1328+
1329+
// TryCast should also strip extension metadata when target has none
1330+
let try_cast_expr = Expr::TryCast(TryCast {
1331+
expr: Box::new(col("foo")),
1332+
field: Arc::new(Field::new("", DataType::Utf8, true)),
1333+
});
1334+
1335+
let (_, result_field) = try_cast_expr.to_field(&schema).unwrap();
1336+
assert!(
1337+
result_field
1338+
.metadata()
1339+
.get(EXTENSION_TYPE_NAME_KEY)
1340+
.is_none(),
1341+
"Extension type name should be stripped in TryCast"
1342+
);
1343+
// TryCast should force nullable
1344+
assert!(
1345+
result_field.is_nullable(),
1346+
"TryCast result should be nullable"
1347+
);
1348+
}
12111349
}

‎datafusion/physical-expr/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ proto = [
5151

5252
[dependencies]
5353
arrow = { workspace = true }
54+
arrow-schema = { workspace = true }
5455
datafusion-common = { workspace = true }
5556
datafusion-expr = { workspace = true }
5657
datafusion-expr-common = { workspace = true }

‎datafusion/physical-expr/src/expressions/cast.rs‎

Lines changed: 106 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ use crate::physical_expr::PhysicalExpr;
2424
use arrow::compute::{CastOptions, can_cast_types};
2525
use arrow::datatypes::{DataType, DataType::*, FieldRef, Schema};
2626
use arrow::record_batch::RecordBatch;
27+
use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY};
2728
use datafusion_common::datatype::DataTypeExt;
2829
use datafusion_common::format::DEFAULT_FORMAT_OPTIONS;
2930
use datafusion_common::nested_struct::{
@@ -151,18 +152,36 @@ impl CastExpr {
151152
}
152153

153154
fn resolved_target_field(&self, input_schema: &Schema) -> Result<FieldRef> {
154-
if is_default_target_field(&self.target_field) {
155-
self.expr.return_field(input_schema).map(|field| {
155+
self.expr.return_field(input_schema).map(|source_field| {
156+
if is_default_target_field(&self.target_field) {
157+
// Type-only cast: derive from source field but strip extension metadata
158+
let mut metadata = source_field.metadata().clone();
159+
metadata.remove(EXTENSION_TYPE_NAME_KEY);
160+
metadata.remove(EXTENSION_TYPE_METADATA_KEY);
161+
156162
Arc::new(
157-
field
163+
source_field
158164
.as_ref()
159165
.clone()
160-
.with_data_type(self.cast_type().clone()),
166+
.with_data_type(self.cast_type().clone())
167+
.with_metadata(metadata),
161168
)
162-
})
163-
} else {
164-
Ok(Arc::clone(&self.target_field))
165-
}
169+
} else {
170+
// Explicit target field: use target's attributes but handle metadata carefully
171+
// - Start with source's non-extension metadata
172+
// - Then add all target metadata (including extension type metadata)
173+
let mut metadata = source_field.metadata().clone();
174+
metadata.remove(EXTENSION_TYPE_NAME_KEY);
175+
metadata.remove(EXTENSION_TYPE_METADATA_KEY);
176+
177+
// Target metadata takes precedence (including extension type metadata)
178+
for (k, v) in self.target_field.metadata() {
179+
metadata.insert(k.clone(), v.clone());
180+
}
181+
182+
Arc::new(self.target_field.as_ref().clone().with_metadata(metadata))
183+
}
184+
})
166185
}
167186

168187
/// Check if casting from the specified source type to the target type is a
@@ -1208,6 +1227,85 @@ mod tests {
12081227

12091228
Ok(())
12101229
}
1230+
1231+
#[test]
1232+
fn type_only_cast_strips_extension_metadata() -> Result<()> {
1233+
// When using type-only cast (new()), extension metadata from source should NOT propagate
1234+
let source_meta = HashMap::from([
1235+
(
1236+
EXTENSION_TYPE_NAME_KEY.to_string(),
1237+
"arrow.uuid".to_string(),
1238+
),
1239+
("custom_key".to_string(), "custom_value".to_string()),
1240+
]);
1241+
let schema = Schema::new(vec![
1242+
Field::new("a", FixedSizeBinary(16), false).with_metadata(source_meta),
1243+
]);
1244+
1245+
let expr = CastExpr::new(col("a", &schema)?, Utf8, None);
1246+
1247+
let field = expr.return_field(&schema)?;
1248+
assert!(
1249+
field.metadata().get(EXTENSION_TYPE_NAME_KEY).is_none(),
1250+
"Type-only cast should strip extension type name from source"
1251+
);
1252+
assert_eq!(
1253+
field.metadata().get("custom_key"),
1254+
Some(&"custom_value".to_string()),
1255+
"Type-only cast should preserve non-extension metadata"
1256+
);
1257+
1258+
Ok(())
1259+
}
1260+
1261+
#[test]
1262+
fn field_aware_cast_applies_target_extension_metadata() -> Result<()> {
1263+
// When using field-aware cast, target's extension metadata should be applied
1264+
let source_meta = HashMap::from([
1265+
(
1266+
EXTENSION_TYPE_NAME_KEY.to_string(),
1267+
"source.type".to_string(),
1268+
),
1269+
("source_key".to_string(), "source_value".to_string()),
1270+
]);
1271+
let target_meta = HashMap::from([
1272+
(
1273+
EXTENSION_TYPE_NAME_KEY.to_string(),
1274+
"target.type".to_string(),
1275+
),
1276+
(
1277+
EXTENSION_TYPE_METADATA_KEY.to_string(),
1278+
"target_ext_meta".to_string(),
1279+
),
1280+
]);
1281+
let schema = Schema::new(vec![
1282+
Field::new("a", FixedSizeBinary(16), false).with_metadata(source_meta),
1283+
]);
1284+
1285+
let target_field =
1286+
Arc::new(Field::new("b", Utf8, true).with_metadata(target_meta));
1287+
let expr =
1288+
CastExpr::new_with_target_field(col("a", &schema)?, target_field, None);
1289+
1290+
let field = expr.return_field(&schema)?;
1291+
assert_eq!(
1292+
field.metadata().get(EXTENSION_TYPE_NAME_KEY),
1293+
Some(&"target.type".to_string()),
1294+
"Field-aware cast should use target's extension type name"
1295+
);
1296+
assert_eq!(
1297+
field.metadata().get(EXTENSION_TYPE_METADATA_KEY),
1298+
Some(&"target_ext_meta".to_string()),
1299+
"Field-aware cast should use target's extension type metadata"
1300+
);
1301+
assert_eq!(
1302+
field.metadata().get("source_key"),
1303+
Some(&"source_value".to_string()),
1304+
"Field-aware cast should preserve source's non-extension metadata"
1305+
);
1306+
1307+
Ok(())
1308+
}
12111309
}
12121310

12131311
/// Tests for the `try_to_proto` / `try_from_proto` hooks.

0 commit comments

Comments
 (0)