Skip to content

Commit d9a9321

Browse files
authored
fix: preserve literal metadata and destructure leaf and unary proto hooks (#25382)
## Which issue does this PR close? Closes #24613. Literal field metadata is lost when physical expressions are serialized. ## Rationale for this change Losing this metadata can strip Arrow extension types and custom metadata from projected columns. ## What changes are included in this PR? Use exhaustive destructuring in all seven leaf and unary expression encoders and decoders. Add a literal protobuf message that carries metadata and regenerate the Rust and JSON bindings. ## What is the testing strategy for this PR? A projection round-trip test covers null and non-null literals through protobuf and JSON and was verified to fail before the fix. Existing tests cover plain literals, and a new test rejects messages missing a value. The extended workspace tests and all-feature Clippy pass. ## Are there any user-facing changes? Literal metadata survives physical plan serialization. Literals with metadata use a new protobuf variant that requires an updated reader.
1 parent bb24345 commit d9a9321

12 files changed

Lines changed: 293 additions & 38 deletions

File tree

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

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -153,26 +153,34 @@ impl PhysicalExpr for Column {
153153
_ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>,
154154
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> {
155155
use datafusion_proto_models::protobuf;
156+
let Self { name, index } = self;
156157
Ok(Some(protobuf::PhysicalExprNode {
157158
expr_id: None,
158-
expr_type: Some(protobuf::physical_expr_node::ExprType::Column(self.into())),
159+
expr_type: Some(protobuf::physical_expr_node::ExprType::Column(
160+
protobuf::PhysicalColumn {
161+
name: name.clone(),
162+
index: *index as u32,
163+
},
164+
)),
159165
}))
160166
}
161167
}
162168

163169
#[cfg(feature = "proto")]
164170
impl From<&datafusion_proto_models::protobuf::PhysicalColumn> for Column {
165171
fn from(c: &datafusion_proto_models::protobuf::PhysicalColumn) -> Self {
166-
Column::new(&c.name, c.index as usize)
172+
let datafusion_proto_models::protobuf::PhysicalColumn { name, index } = c;
173+
Column::new(name, *index as usize)
167174
}
168175
}
169176

170177
#[cfg(feature = "proto")]
171178
impl From<&Column> for datafusion_proto_models::protobuf::PhysicalColumn {
172179
fn from(c: &Column) -> Self {
180+
let Column { name, index } = c;
173181
Self {
174-
name: c.name.clone(),
175-
index: c.index as u32,
182+
name: name.clone(),
183+
index: *index as u32,
176184
}
177185
}
178186
}
@@ -196,12 +204,12 @@ impl Column {
196204
) -> Result<Arc<dyn PhysicalExpr>> {
197205
use datafusion_physical_expr_common::expect_expr_variant;
198206
use datafusion_proto_models::protobuf;
199-
let column = expect_expr_variant!(
207+
let protobuf::PhysicalColumn { name, index } = expect_expr_variant!(
200208
node,
201209
protobuf::physical_expr_node::ExprType::Column,
202210
"Column",
203211
);
204-
Ok(Arc::new(Column::from(column)))
212+
Ok(Arc::new(Column::new(name, *index as usize)))
205213
}
206214
}
207215

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

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -110,11 +110,12 @@ impl PhysicalExpr for IsNotNullExpr {
110110
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> {
111111
use datafusion_proto_models::protobuf;
112112

113+
let Self { arg } = self;
113114
Ok(Some(protobuf::PhysicalExprNode {
114115
expr_id: None,
115116
expr_type: Some(protobuf::physical_expr_node::ExprType::IsNotNullExpr(
116117
Box::new(protobuf::PhysicalIsNotNull {
117-
expr: Some(Box::new(ctx.encode_child(&self.arg)?)),
118+
expr: Some(Box::new(ctx.encode_child(arg)?)),
118119
}),
119120
)),
120121
}))
@@ -136,11 +137,9 @@ impl IsNotNullExpr {
136137
protobuf::physical_expr_node::ExprType::IsNotNullExpr,
137138
"IsNotNullExpr",
138139
);
139-
let expr = ctx.decode_required_expression(
140-
node.expr.as_deref(),
141-
"IsNotNullExpr",
142-
"expr",
143-
)?;
140+
let protobuf::PhysicalIsNotNull { expr } = node.as_ref();
141+
let expr =
142+
ctx.decode_required_expression(expr.as_deref(), "IsNotNullExpr", "expr")?;
144143

145144
Ok(Arc::new(IsNotNullExpr::new(expr)))
146145
}

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -109,11 +109,12 @@ impl PhysicalExpr for IsNullExpr {
109109
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> {
110110
use datafusion_proto_models::protobuf;
111111

112+
let Self { arg } = self;
112113
Ok(Some(protobuf::PhysicalExprNode {
113114
expr_id: None,
114115
expr_type: Some(protobuf::physical_expr_node::ExprType::IsNullExpr(
115116
Box::new(protobuf::PhysicalIsNull {
116-
expr: Some(Box::new(ctx.encode_child(&self.arg)?)),
117+
expr: Some(Box::new(ctx.encode_child(arg)?)),
117118
}),
118119
)),
119120
}))
@@ -135,8 +136,9 @@ impl IsNullExpr {
135136
protobuf::physical_expr_node::ExprType::IsNullExpr,
136137
"IsNullExpr",
137138
);
139+
let protobuf::PhysicalIsNull { expr } = node.as_ref();
138140
let expr =
139-
ctx.decode_required_expression(node.expr.as_deref(), "IsNullExpr", "expr")?;
141+
ctx.decode_required_expression(expr.as_deref(), "IsNullExpr", "expr")?;
140142

141143
Ok(Arc::new(IsNullExpr::new(expr)))
142144
}

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

Lines changed: 64 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -143,11 +143,21 @@ impl PhysicalExpr for Literal {
143143
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> {
144144
use datafusion_proto_models::protobuf;
145145

146+
let Self { value, field } = self;
147+
// The field name, type, and nullability are reconstructed by new_with_metadata.
148+
let expr_type = if field.metadata().is_empty() {
149+
protobuf::physical_expr_node::ExprType::Literal(value.try_into()?)
150+
} else {
151+
protobuf::physical_expr_node::ExprType::LiteralWithMetadata(
152+
protobuf::PhysicalLiteralNode {
153+
value: Some(value.try_into()?),
154+
metadata: field.metadata().into(),
155+
},
156+
)
157+
};
146158
Ok(Some(protobuf::PhysicalExprNode {
147159
expr_id: None,
148-
expr_type: Some(protobuf::physical_expr_node::ExprType::Literal(
149-
(&self.value).try_into()?,
150-
)),
160+
expr_type: Some(expr_type),
151161
}))
152162
}
153163
}
@@ -159,16 +169,38 @@ impl Literal {
159169
node: &datafusion_proto_models::protobuf::PhysicalExprNode,
160170
_ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>,
161171
) -> Result<Arc<dyn PhysicalExpr>> {
162-
use datafusion_physical_expr_common::expect_expr_variant;
172+
use datafusion_common::{internal_datafusion_err, internal_err};
163173
use datafusion_proto_models::protobuf;
164-
165-
let scalar_proto = expect_expr_variant!(
166-
node,
167-
protobuf::physical_expr_node::ExprType::Literal,
168-
"Literal",
169-
);
170-
let value = ScalarValue::try_from(scalar_proto)?;
171-
Ok(Arc::new(Literal::new(value)))
174+
use protobuf::physical_expr_node::ExprType;
175+
176+
let protobuf::PhysicalExprNode {
177+
expr_type,
178+
// Expression IDs are handled by the enclosing proto converter.
179+
expr_id: _,
180+
} = node;
181+
let (value, metadata) = match expr_type {
182+
Some(ExprType::Literal(scalar)) => {
183+
let datafusion_proto_models::datafusion_common::ScalarValue {
184+
// The scalar payload is decoded by ScalarValue::try_from.
185+
value: _,
186+
} = scalar;
187+
(ScalarValue::try_from(scalar)?, None)
188+
}
189+
Some(ExprType::LiteralWithMetadata(protobuf::PhysicalLiteralNode {
190+
value,
191+
metadata,
192+
})) => {
193+
let value = value.as_ref().ok_or_else(|| {
194+
internal_datafusion_err!("Literal is missing required field 'value'")
195+
})?;
196+
(
197+
ScalarValue::try_from(value)?,
198+
Some(FieldMetadata::from(metadata)),
199+
)
200+
}
201+
_ => return internal_err!("PhysicalExprNode is not a Literal"),
202+
};
203+
Ok(Arc::new(Literal::new_with_metadata(value, metadata)))
172204
}
173205
}
174206

@@ -315,6 +347,26 @@ mod proto_tests {
315347
assert_eq!(lit.value(), &ScalarValue::Int32(Some(42)));
316348
}
317349

350+
#[test]
351+
fn try_from_proto_rejects_missing_value() {
352+
let node = datafusion_proto_models::protobuf::PhysicalExprNode {
353+
expr_id: None,
354+
expr_type: Some(physical_expr_node::ExprType::LiteralWithMetadata(
355+
datafusion_proto_models::protobuf::PhysicalLiteralNode {
356+
value: None,
357+
metadata: Default::default(),
358+
},
359+
)),
360+
};
361+
let schema = Schema::empty();
362+
let decoder = UnreachableDecoder;
363+
let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
364+
let err = Literal::try_from_proto(&node, &ctx).unwrap_err();
365+
assert!(
366+
matches!(err, DataFusionError::Internal(msg) if msg.contains("Literal is missing required field 'value'"))
367+
);
368+
}
369+
318370
#[test]
319371
fn try_from_proto_rejects_non_literal_node() {
320372
let node = column_node("a");

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -184,11 +184,12 @@ impl PhysicalExpr for NegativeExpr {
184184
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> {
185185
use datafusion_proto_models::protobuf;
186186

187+
let Self { arg } = self;
187188
Ok(Some(protobuf::PhysicalExprNode {
188189
expr_id: None,
189190
expr_type: Some(protobuf::physical_expr_node::ExprType::Negative(Box::new(
190191
protobuf::PhysicalNegativeNode {
191-
expr: Some(Box::new(ctx.encode_child(&self.arg)?)),
192+
expr: Some(Box::new(ctx.encode_child(arg)?)),
192193
},
193194
))),
194195
}))
@@ -210,8 +211,9 @@ impl NegativeExpr {
210211
protobuf::physical_expr_node::ExprType::Negative,
211212
"Negative",
212213
);
214+
let protobuf::PhysicalNegativeNode { expr } = n.as_ref();
213215
let expr =
214-
ctx.decode_required_expression(n.expr.as_deref(), "NegativeExpr", "expr")?;
216+
ctx.decode_required_expression(expr.as_deref(), "NegativeExpr", "expr")?;
215217

216218
Ok(Arc::new(NegativeExpr::new(expr)))
217219
}

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -189,11 +189,12 @@ impl PhysicalExpr for NotExpr {
189189
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> {
190190
use datafusion_proto_models::protobuf;
191191

192+
let Self { arg } = self;
192193
Ok(Some(protobuf::PhysicalExprNode {
193194
expr_id: None,
194195
expr_type: Some(protobuf::physical_expr_node::ExprType::NotExpr(Box::new(
195196
protobuf::PhysicalNot {
196-
expr: Some(Box::new(ctx.encode_child(&self.arg)?)),
197+
expr: Some(Box::new(ctx.encode_child(arg)?)),
197198
},
198199
))),
199200
}))
@@ -215,8 +216,8 @@ impl NotExpr {
215216
protobuf::physical_expr_node::ExprType::NotExpr,
216217
"NotExpr",
217218
);
218-
let expr =
219-
ctx.decode_required_expression(not_expr.expr.as_deref(), "NotExpr", "expr")?;
219+
let protobuf::PhysicalNot { expr } = not_expr.as_ref();
220+
let expr = ctx.decode_required_expression(expr.as_deref(), "NotExpr", "expr")?;
220221

221222
Ok(Arc::new(NotExpr::new(expr)))
222223
}

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

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -93,12 +93,11 @@ impl PhysicalExpr for UnKnownColumn {
9393
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> {
9494
use datafusion_proto_models::protobuf;
9595

96+
let Self { name } = self;
9697
Ok(Some(protobuf::PhysicalExprNode {
9798
expr_id: None,
9899
expr_type: Some(protobuf::physical_expr_node::ExprType::UnknownColumn(
99-
protobuf::UnknownColumn {
100-
name: self.name.clone(),
101-
},
100+
protobuf::UnknownColumn { name: name.clone() },
102101
)),
103102
}))
104103
}
@@ -114,12 +113,12 @@ impl UnKnownColumn {
114113
use datafusion_physical_expr_common::expect_expr_variant;
115114
use datafusion_proto_models::protobuf;
116115

117-
let unknown_col = expect_expr_variant!(
116+
let protobuf::UnknownColumn { name } = expect_expr_variant!(
118117
node,
119118
protobuf::physical_expr_node::ExprType::UnknownColumn,
120119
"UnKnownColumn",
121120
);
122-
Ok(Arc::new(UnKnownColumn::new(&unknown_col.name)))
121+
Ok(Arc::new(UnKnownColumn::new(name)))
123122
}
124123
}
125124

‎datafusion/proto-models/proto/datafusion.proto‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1105,9 +1105,15 @@ message PhysicalExprNode {
11051105
PhysicalLambdaVariableExprNode lambda_variable = 26;
11061106
PhysicalRangeExprNode range_expr = 27;
11071107
PhysicalSqlSimilarToPatternNode sql_similar_to_pattern = 28;
1108+
PhysicalLiteralNode literal_with_metadata = 29;
11081109
}
11091110
}
11101111

1112+
message PhysicalLiteralNode {
1113+
datafusion_common.ScalarValue value = 1;
1114+
map<string, string> metadata = 2;
1115+
}
1116+
11111117
message PhysicalDynamicFilterNode {
11121118
repeated PhysicalExprNode children = 1;
11131119
repeated PhysicalExprNode remapped_children = 2;

0 commit comments

Comments
 (0)