You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by bl...@apache.org on 2026/02/06 14:58:50 UTC
(datafusion-sandbox) 01/44: Preserve PhysicalExpr graph in proto round trip using Arc pointers as unique identifiers (#20037)
This is an automated email from the ASF dual-hosted git repository.
blaginin pushed a commit to branch annarose/dict-coercion
in repository https://gitbox.apache.org/repos/asf/datafusion-sandbox.git
commit 66ee0afcabfd59323eeb7b5f76f60a80881bbba3
Author: Adrian Garcia Badaracco <17...@users.noreply.github.com>
AuthorDate: Mon Feb 2 08:04:11 2026 -0600
Preserve PhysicalExpr graph in proto round trip using Arc pointers as unique identifiers (#20037)
Replaces #18192 using the APIs in #19437.
Similar to #18192 the end goal here is specifically to enable
deduplication of `DynamicFilterPhysicalExpr` so that distributed query
engines can get one step closer to using dynamic filters.
Because it's actually simpler we apply this deduplication to all
`PhysicalExpr`s with the added benefit that we more faithfully preserve
the original expression tree (instead of adding new duplicate branches)
which will have the immediate impact of e.g. not duplicating large
`InListExpr`s.
---
Cargo.lock | 1 +
datafusion/proto/Cargo.toml | 1 +
datafusion/proto/proto/datafusion.proto | 8 +
datafusion/proto/src/generated/pbjson.rs | 22 +
datafusion/proto/src/generated/prost.rs | 8 +
datafusion/proto/src/physical_plan/mod.rs | 216 ++++++++++
datafusion/proto/src/physical_plan/to_proto.rs | 18 +
.../proto/tests/cases/roundtrip_physical_plan.rs | 452 ++++++++++++++++++++-
8 files changed, 724 insertions(+), 2 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index 1f28687f4..e582b43c7 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2514,6 +2514,7 @@ dependencies = [
"pbjson 0.9.0",
"pretty_assertions",
"prost",
+ "rand 0.9.2",
"serde",
"serde_json",
"tokio",
diff --git a/datafusion/proto/Cargo.toml b/datafusion/proto/Cargo.toml
index edc96d5a9..3d17ed30d 100644
--- a/datafusion/proto/Cargo.toml
+++ b/datafusion/proto/Cargo.toml
@@ -66,6 +66,7 @@ datafusion-proto-common = { workspace = true }
object_store = { workspace = true }
pbjson = { workspace = true, optional = true }
prost = { workspace = true }
+rand = { workspace = true }
serde = { version = "1.0", optional = true }
serde_json = { workspace = true, optional = true }
diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto
index 59be5c578..2c29597f4 100644
--- a/datafusion/proto/proto/datafusion.proto
+++ b/datafusion/proto/proto/datafusion.proto
@@ -850,6 +850,14 @@ message PhysicalExprNode {
// Was date_time_interval_expr
reserved 17;
+ // Unique identifier for this expression to do deduplication during deserialization.
+ // When serializing, this is set to a unique identifier for each combination of
+ // expression, process and serialization run.
+ // When deserializing, if this ID has been seen before, the cached Arc is returned
+ // instead of creating a new one, enabling reconstruction of referential integrity
+ // across serde roundtrips.
+ optional uint64 expr_id = 30;
+
oneof ExprType {
// column references
PhysicalColumn column = 1;
diff --git a/datafusion/proto/src/generated/pbjson.rs b/datafusion/proto/src/generated/pbjson.rs
index 3873afcdc..5d8ee5b73 100644
--- a/datafusion/proto/src/generated/pbjson.rs
+++ b/datafusion/proto/src/generated/pbjson.rs
@@ -16029,10 +16029,18 @@ impl serde::Serialize for PhysicalExprNode {
{
use serde::ser::SerializeStruct;
let mut len = 0;
+ if self.expr_id.is_some() {
+ len += 1;
+ }
if self.expr_type.is_some() {
len += 1;
}
let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalExprNode", len)?;
+ if let Some(v) = self.expr_id.as_ref() {
+ #[allow(clippy::needless_borrow)]
+ #[allow(clippy::needless_borrows_for_generic_args)]
+ struct_ser.serialize_field("exprId", ToString::to_string(&v).as_str())?;
+ }
if let Some(v) = self.expr_type.as_ref() {
match v {
physical_expr_node::ExprType::Column(v) => {
@@ -16104,6 +16112,8 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode {
D: serde::Deserializer<'de>,
{
const FIELDS: &[&str] = &[
+ "expr_id",
+ "exprId",
"column",
"literal",
"binary_expr",
@@ -16140,6 +16150,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode {
#[allow(clippy::enum_variant_names)]
enum GeneratedField {
+ ExprId,
Column,
Literal,
BinaryExpr,
@@ -16180,6 +16191,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode {
E: serde::de::Error,
{
match value {
+ "exprId" | "expr_id" => Ok(GeneratedField::ExprId),
"column" => Ok(GeneratedField::Column),
"literal" => Ok(GeneratedField::Literal),
"binaryExpr" | "binary_expr" => Ok(GeneratedField::BinaryExpr),
@@ -16218,9 +16230,18 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode {
where
V: serde::de::MapAccess<'de>,
{
+ let mut expr_id__ = None;
let mut expr_type__ = None;
while let Some(k) = map_.next_key()? {
match k {
+ GeneratedField::ExprId => {
+ if expr_id__.is_some() {
+ return Err(serde::de::Error::duplicate_field("exprId"));
+ }
+ expr_id__ =
+ map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| x.0)
+ ;
+ }
GeneratedField::Column => {
if expr_type__.is_some() {
return Err(serde::de::Error::duplicate_field("column"));
@@ -16357,6 +16378,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode {
}
}
Ok(PhysicalExprNode {
+ expr_id: expr_id__,
expr_type: expr_type__,
})
}
diff --git a/datafusion/proto/src/generated/prost.rs b/datafusion/proto/src/generated/prost.rs
index 3806e31a4..18dabac51 100644
--- a/datafusion/proto/src/generated/prost.rs
+++ b/datafusion/proto/src/generated/prost.rs
@@ -1280,6 +1280,14 @@ pub struct PhysicalExtensionNode {
/// physical expressions
#[derive(Clone, PartialEq, ::prost::Message)]
pub struct PhysicalExprNode {
+ /// Unique identifier for this expression to do deduplication during deserialization.
+ /// When serializing, this is set to a unique identifier for each combination of
+ /// expression, process and serialization run.
+ /// When deserializing, if this ID has been seen before, the cached Arc is returned
+ /// instead of creating a new one, enabling reconstruction of referential integrity
+ /// across serde roundtrips.
+ #[prost(uint64, optional, tag = "30")]
+ pub expr_id: ::core::option::Option<u64>,
#[prost(
oneof = "physical_expr_node::ExprType",
tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 15, 16, 18, 19, 20, 21"
diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs
index e1f6381d1..2b805c3a2 100644
--- a/datafusion/proto/src/physical_plan/mod.rs
+++ b/datafusion/proto/src/physical_plan/mod.rs
@@ -15,7 +15,10 @@
// specific language governing permissions and limitations
// under the License.
+use std::cell::RefCell;
+use std::collections::HashMap;
use std::fmt::Debug;
+use std::hash::{DefaultHasher, Hash, Hasher};
use std::sync::Arc;
use arrow::compute::SortOptions;
@@ -2993,6 +2996,7 @@ impl protobuf::PhysicalPlanNode {
nulls_first: expr.options.nulls_first,
});
Ok(protobuf::PhysicalExprNode {
+ expr_id: None,
expr_type: Some(ExprType::Sort(sort_expr)),
})
})
@@ -3078,6 +3082,7 @@ impl protobuf::PhysicalPlanNode {
nulls_first: expr.options.nulls_first,
});
Ok(protobuf::PhysicalExprNode {
+ expr_id: None,
expr_type: Some(ExprType::Sort(sort_expr)),
})
})
@@ -3712,6 +3717,217 @@ impl PhysicalProtoConverterExtension for DefaultPhysicalProtoConverter {
}
}
+/// Internal serializer that adds expr_id to expressions.
+/// Created fresh for each serialization operation.
+struct DeduplicatingSerializer {
+ /// Random salt combined with pointer addresses and process ID to create globally unique expr_ids.
+ session_id: u64,
+}
+
+impl DeduplicatingSerializer {
+ fn new() -> Self {
+ Self {
+ session_id: rand::random(),
+ }
+ }
+}
+
+impl PhysicalProtoConverterExtension for DeduplicatingSerializer {
+ fn proto_to_execution_plan(
+ &self,
+ _ctx: &TaskContext,
+ _codec: &dyn PhysicalExtensionCodec,
+ _proto: &protobuf::PhysicalPlanNode,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ internal_err!("DeduplicatingSerializer cannot deserialize execution plans")
+ }
+
+ fn execution_plan_to_proto(
+ &self,
+ plan: &Arc<dyn ExecutionPlan>,
+ codec: &dyn PhysicalExtensionCodec,
+ ) -> Result<protobuf::PhysicalPlanNode>
+ where
+ Self: Sized,
+ {
+ protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter(
+ Arc::clone(plan),
+ codec,
+ self,
+ )
+ }
+
+ fn proto_to_physical_expr(
+ &self,
+ _proto: &protobuf::PhysicalExprNode,
+ _ctx: &TaskContext,
+ _input_schema: &Schema,
+ _codec: &dyn PhysicalExtensionCodec,
+ ) -> Result<Arc<dyn PhysicalExpr>>
+ where
+ Self: Sized,
+ {
+ internal_err!("DeduplicatingSerializer cannot deserialize physical expressions")
+ }
+
+ fn physical_expr_to_proto(
+ &self,
+ expr: &Arc<dyn PhysicalExpr>,
+ codec: &dyn PhysicalExtensionCodec,
+ ) -> Result<protobuf::PhysicalExprNode> {
+ let mut proto = serialize_physical_expr_with_converter(expr, codec, self)?;
+
+ // Hash session_id, pointer address, and process ID together to create expr_id.
+ // - session_id: random per serializer, prevents collisions when merging serializations
+ // - ptr: unique address per Arc within a process
+ // - pid: prevents collisions if serializer is shared across processes
+ let mut hasher = DefaultHasher::new();
+ self.session_id.hash(&mut hasher);
+ (Arc::as_ptr(expr) as *const () as u64).hash(&mut hasher);
+ std::process::id().hash(&mut hasher);
+ proto.expr_id = Some(hasher.finish());
+
+ Ok(proto)
+ }
+}
+
+/// Internal deserializer that caches expressions by expr_id.
+/// Created fresh for each deserialization operation.
+#[derive(Default)]
+struct DeduplicatingDeserializer {
+ /// Cache mapping expr_id to deserialized expressions.
+ cache: RefCell<HashMap<u64, Arc<dyn PhysicalExpr>>>,
+}
+
+impl PhysicalProtoConverterExtension for DeduplicatingDeserializer {
+ fn proto_to_execution_plan(
+ &self,
+ ctx: &TaskContext,
+ codec: &dyn PhysicalExtensionCodec,
+ proto: &protobuf::PhysicalPlanNode,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ proto.try_into_physical_plan_with_converter(ctx, codec, self)
+ }
+
+ fn execution_plan_to_proto(
+ &self,
+ _plan: &Arc<dyn ExecutionPlan>,
+ _codec: &dyn PhysicalExtensionCodec,
+ ) -> Result<protobuf::PhysicalPlanNode>
+ where
+ Self: Sized,
+ {
+ internal_err!("DeduplicatingDeserializer cannot serialize execution plans")
+ }
+
+ fn proto_to_physical_expr(
+ &self,
+ proto: &protobuf::PhysicalExprNode,
+ ctx: &TaskContext,
+ input_schema: &Schema,
+ codec: &dyn PhysicalExtensionCodec,
+ ) -> Result<Arc<dyn PhysicalExpr>>
+ where
+ Self: Sized,
+ {
+ if let Some(expr_id) = proto.expr_id {
+ // Check cache first
+ if let Some(cached) = self.cache.borrow().get(&expr_id) {
+ return Ok(Arc::clone(cached));
+ }
+ // Deserialize and cache
+ let expr = parse_physical_expr_with_converter(
+ proto,
+ ctx,
+ input_schema,
+ codec,
+ self,
+ )?;
+ self.cache.borrow_mut().insert(expr_id, Arc::clone(&expr));
+ Ok(expr)
+ } else {
+ parse_physical_expr_with_converter(proto, ctx, input_schema, codec, self)
+ }
+ }
+
+ fn physical_expr_to_proto(
+ &self,
+ _expr: &Arc<dyn PhysicalExpr>,
+ _codec: &dyn PhysicalExtensionCodec,
+ ) -> Result<protobuf::PhysicalExprNode> {
+ internal_err!("DeduplicatingDeserializer cannot serialize physical expressions")
+ }
+}
+
+/// A proto converter that adds expression deduplication during serialization
+/// and deserialization.
+///
+/// During serialization, each expression's Arc pointer address is XORed with a
+/// random session_id to create a salted `expr_id`. This prevents cross-process
+/// collisions when serialized plans are merged.
+///
+/// During deserialization, expressions with the same `expr_id` share the same
+/// Arc, reducing memory usage for plans with duplicate expressions (e.g., large
+/// IN lists) and supporting correctly linking [`DynamicFilterPhysicalExpr`] instances.
+///
+/// This converter is stateless - it creates internal serializers/deserializers
+/// on demand for each operation.
+///
+/// [`DynamicFilterPhysicalExpr`]: https://docs.rs/datafusion-physical-expr/latest/datafusion_physical_expr/expressions/struct.DynamicFilterPhysicalExpr.html
+#[derive(Debug, Default, Clone, Copy)]
+pub struct DeduplicatingProtoConverter {}
+
+impl PhysicalProtoConverterExtension for DeduplicatingProtoConverter {
+ fn proto_to_execution_plan(
+ &self,
+ ctx: &TaskContext,
+ codec: &dyn PhysicalExtensionCodec,
+ proto: &protobuf::PhysicalPlanNode,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let deserializer = DeduplicatingDeserializer::default();
+ proto.try_into_physical_plan_with_converter(ctx, codec, &deserializer)
+ }
+
+ fn execution_plan_to_proto(
+ &self,
+ plan: &Arc<dyn ExecutionPlan>,
+ codec: &dyn PhysicalExtensionCodec,
+ ) -> Result<protobuf::PhysicalPlanNode>
+ where
+ Self: Sized,
+ {
+ let serializer = DeduplicatingSerializer::new();
+ protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter(
+ Arc::clone(plan),
+ codec,
+ &serializer,
+ )
+ }
+
+ fn proto_to_physical_expr(
+ &self,
+ proto: &protobuf::PhysicalExprNode,
+ ctx: &TaskContext,
+ input_schema: &Schema,
+ codec: &dyn PhysicalExtensionCodec,
+ ) -> Result<Arc<dyn PhysicalExpr>>
+ where
+ Self: Sized,
+ {
+ let deserializer = DeduplicatingDeserializer::default();
+ deserializer.proto_to_physical_expr(proto, ctx, input_schema, codec)
+ }
+
+ fn physical_expr_to_proto(
+ &self,
+ expr: &Arc<dyn PhysicalExpr>,
+ codec: &dyn PhysicalExtensionCodec,
+ ) -> Result<protobuf::PhysicalExprNode> {
+ let serializer = DeduplicatingSerializer::new();
+ serializer.physical_expr_to_proto(expr, codec)
+ }
+}
+
/// A PhysicalExtensionCodec that tries one of multiple inner codecs
/// until one works
#[derive(Debug)]
diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs
index 08ce00da4..a38e59acd 100644
--- a/datafusion/proto/src/physical_plan/to_proto.rs
+++ b/datafusion/proto/src/physical_plan/to_proto.rs
@@ -71,6 +71,7 @@ pub fn serialize_physical_aggr_expr(
let mut buf = Vec::new();
codec.try_encode_udaf(aggr_expr.fun(), &mut buf)?;
Ok(protobuf::PhysicalExprNode {
+ expr_id: None,
expr_type: Some(protobuf::physical_expr_node::ExprType::AggregateExpr(
protobuf::PhysicalAggregateExprNode {
aggregate_function: Some(physical_aggregate_expr_node::AggregateFunction::UserDefinedAggrFunction(name)),
@@ -280,12 +281,14 @@ pub fn serialize_physical_expr_with_converter(
)),
};
return Ok(protobuf::PhysicalExprNode {
+ expr_id: None,
expr_type: Some(protobuf::physical_expr_node::ExprType::Literal(value)),
});
}
if let Some(expr) = expr.downcast_ref::<Column>() {
Ok(protobuf::PhysicalExprNode {
+ expr_id: None,
expr_type: Some(protobuf::physical_expr_node::ExprType::Column(
protobuf::PhysicalColumn {
name: expr.name().to_string(),
@@ -295,6 +298,7 @@ pub fn serialize_physical_expr_with_converter(
})
} else if let Some(expr) = expr.downcast_ref::<UnKnownColumn>() {
Ok(protobuf::PhysicalExprNode {
+ expr_id: None,
expr_type: Some(protobuf::physical_expr_node::ExprType::UnknownColumn(
protobuf::UnknownColumn {
name: expr.name().to_string(),
@@ -313,12 +317,14 @@ pub fn serialize_physical_expr_with_converter(
});
Ok(protobuf::PhysicalExprNode {
+ expr_id: None,
expr_type: Some(protobuf::physical_expr_node::ExprType::BinaryExpr(
binary_expr,
)),
})
} else if let Some(expr) = expr.downcast_ref::<CaseExpr>() {
Ok(protobuf::PhysicalExprNode {
+ expr_id: None,
expr_type: Some(
protobuf::physical_expr_node::ExprType::Case(
Box::new(
@@ -361,6 +367,7 @@ pub fn serialize_physical_expr_with_converter(
})
} else if let Some(expr) = expr.downcast_ref::<NotExpr>() {
Ok(protobuf::PhysicalExprNode {
+ expr_id: None,
expr_type: Some(protobuf::physical_expr_node::ExprType::NotExpr(Box::new(
protobuf::PhysicalNot {
expr: Some(Box::new(
@@ -371,6 +378,7 @@ pub fn serialize_physical_expr_with_converter(
})
} else if let Some(expr) = expr.downcast_ref::<IsNullExpr>() {
Ok(protobuf::PhysicalExprNode {
+ expr_id: None,
expr_type: Some(protobuf::physical_expr_node::ExprType::IsNullExpr(
Box::new(protobuf::PhysicalIsNull {
expr: Some(Box::new(
@@ -381,6 +389,7 @@ pub fn serialize_physical_expr_with_converter(
})
} else if let Some(expr) = expr.downcast_ref::<IsNotNullExpr>() {
Ok(protobuf::PhysicalExprNode {
+ expr_id: None,
expr_type: Some(protobuf::physical_expr_node::ExprType::IsNotNullExpr(
Box::new(protobuf::PhysicalIsNotNull {
expr: Some(Box::new(
@@ -391,6 +400,7 @@ pub fn serialize_physical_expr_with_converter(
})
} else if let Some(expr) = expr.downcast_ref::<InListExpr>() {
Ok(protobuf::PhysicalExprNode {
+ expr_id: None,
expr_type: Some(protobuf::physical_expr_node::ExprType::InList(Box::new(
protobuf::PhysicalInListNode {
expr: Some(Box::new(
@@ -403,6 +413,7 @@ pub fn serialize_physical_expr_with_converter(
})
} else if let Some(expr) = expr.downcast_ref::<NegativeExpr>() {
Ok(protobuf::PhysicalExprNode {
+ expr_id: None,
expr_type: Some(protobuf::physical_expr_node::ExprType::Negative(Box::new(
protobuf::PhysicalNegativeNode {
expr: Some(Box::new(
@@ -413,12 +424,14 @@ pub fn serialize_physical_expr_with_converter(
})
} else if let Some(lit) = expr.downcast_ref::<Literal>() {
Ok(protobuf::PhysicalExprNode {
+ expr_id: None,
expr_type: Some(protobuf::physical_expr_node::ExprType::Literal(
lit.value().try_into()?,
)),
})
} else if let Some(cast) = expr.downcast_ref::<CastExpr>() {
Ok(protobuf::PhysicalExprNode {
+ expr_id: None,
expr_type: Some(protobuf::physical_expr_node::ExprType::Cast(Box::new(
protobuf::PhysicalCastNode {
expr: Some(Box::new(
@@ -430,6 +443,7 @@ pub fn serialize_physical_expr_with_converter(
})
} else if let Some(cast) = expr.downcast_ref::<TryCastExpr>() {
Ok(protobuf::PhysicalExprNode {
+ expr_id: None,
expr_type: Some(protobuf::physical_expr_node::ExprType::TryCast(Box::new(
protobuf::PhysicalTryCastNode {
expr: Some(Box::new(
@@ -443,6 +457,7 @@ pub fn serialize_physical_expr_with_converter(
let mut buf = Vec::new();
codec.try_encode_udf(expr.fun(), &mut buf)?;
Ok(protobuf::PhysicalExprNode {
+ expr_id: None,
expr_type: Some(protobuf::physical_expr_node::ExprType::ScalarUdf(
protobuf::PhysicalScalarUdfNode {
name: expr.name().to_string(),
@@ -459,6 +474,7 @@ pub fn serialize_physical_expr_with_converter(
})
} else if let Some(expr) = expr.downcast_ref::<LikeExpr>() {
Ok(protobuf::PhysicalExprNode {
+ expr_id: None,
expr_type: Some(protobuf::physical_expr_node::ExprType::LikeExpr(Box::new(
protobuf::PhysicalLikeExprNode {
negated: expr.negated(),
@@ -475,6 +491,7 @@ pub fn serialize_physical_expr_with_converter(
} else if let Some(expr) = expr.downcast_ref::<HashExpr>() {
let (s0, s1, s2, s3) = expr.seeds();
Ok(protobuf::PhysicalExprNode {
+ expr_id: None,
expr_type: Some(protobuf::physical_expr_node::ExprType::HashExpr(
protobuf::PhysicalHashExprNode {
on_columns: serialize_physical_exprs(
@@ -500,6 +517,7 @@ pub fn serialize_physical_expr_with_converter(
.map(|e| proto_converter.physical_expr_to_proto(e, codec))
.collect::<Result<_>>()?;
Ok(protobuf::PhysicalExprNode {
+ expr_id: None,
expr_type: Some(protobuf::physical_expr_node::ExprType::Extension(
protobuf::PhysicalExtensionExprNode { expr: buf, inputs },
)),
diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
index 5bb771137..0a26025a3 100644
--- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
+++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
@@ -116,8 +116,9 @@ use datafusion_proto::bytes::{
use datafusion_proto::physical_plan::from_proto::parse_physical_expr_with_converter;
use datafusion_proto::physical_plan::to_proto::serialize_physical_expr_with_converter;
use datafusion_proto::physical_plan::{
- AsExecutionPlan, DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter,
- PhysicalExtensionCodec, PhysicalProtoConverterExtension,
+ AsExecutionPlan, DeduplicatingProtoConverter, DefaultPhysicalExtensionCodec,
+ DefaultPhysicalProtoConverter, PhysicalExtensionCodec,
+ PhysicalProtoConverterExtension,
};
use datafusion_proto::protobuf;
use datafusion_proto::protobuf::{PhysicalExprNode, PhysicalPlanNode};
@@ -2564,3 +2565,450 @@ fn custom_proto_converter_intercepts() -> Result<()> {
Ok(())
}
+
+/// Test that expression deduplication works during deserialization.
+/// When the same expression Arc is serialized multiple times, it should be
+/// deduplicated on deserialization (sharing the same Arc).
+#[test]
+fn test_expression_deduplication() -> Result<()> {
+ let field_a = Field::new("a", DataType::Int64, false);
+ let schema = Arc::new(Schema::new(vec![field_a]));
+
+ // Create a shared expression that will be used multiple times
+ let shared_col: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
+
+ // Create an InList expression that uses the same column Arc multiple times
+ // This simulates a real-world scenario where expressions are shared
+ let in_list_expr = in_list(
+ Arc::clone(&shared_col),
+ vec![lit(1i64), lit(2i64), lit(3i64)],
+ &false,
+ &schema,
+ )?;
+
+ // Create a binary expression that uses the shared column and the in_list result
+ let binary_expr: Arc<dyn PhysicalExpr> = Arc::new(BinaryExpr::new(
+ Arc::clone(&shared_col),
+ Operator::Eq,
+ lit(42i64),
+ ));
+
+ // Create a plan that has both expressions (they share the `shared_col` Arc)
+ let input = Arc::new(EmptyExec::new(schema.clone()));
+ let filter = FilterExecBuilder::new(in_list_expr, input).build()?;
+ let projection_exprs = vec![ProjectionExpr {
+ expr: binary_expr,
+ alias: "result".to_string(),
+ }];
+ let exec_plan =
+ Arc::new(ProjectionExec::try_new(projection_exprs, Arc::new(filter))?);
+
+ let ctx = SessionContext::new();
+ let codec = DefaultPhysicalExtensionCodec {};
+ let proto_converter = DeduplicatingProtoConverter {};
+
+ // Perform roundtrip
+ let bytes = physical_plan_to_bytes_with_proto_converter(
+ Arc::clone(&exec_plan) as Arc<dyn ExecutionPlan>,
+ &codec,
+ &proto_converter,
+ )?;
+
+ // Create a new converter for deserialization (fresh cache)
+ let deser_converter = DeduplicatingProtoConverter {};
+ let result_plan = physical_plan_from_bytes_with_proto_converter(
+ bytes.as_ref(),
+ ctx.task_ctx().as_ref(),
+ &codec,
+ &deser_converter,
+ )?;
+
+ // Verify the plan structure is correct
+ pretty_assertions::assert_eq!(format!("{exec_plan:?}"), format!("{result_plan:?}"));
+
+ Ok(())
+}
+
+/// Test that expression deduplication correctly shares Arcs for identical expressions.
+/// This test verifies the core deduplication behavior.
+#[test]
+fn test_expression_deduplication_arc_sharing() -> Result<()> {
+ use datafusion_proto::bytes::{
+ physical_plan_from_bytes_with_proto_converter,
+ physical_plan_to_bytes_with_proto_converter,
+ };
+
+ let field_a = Field::new("a", DataType::Int64, false);
+ let schema = Arc::new(Schema::new(vec![field_a]));
+
+ // Create a column expression
+ let col_expr: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
+
+ // Create a projection that uses the SAME Arc twice
+ // After roundtrip, both should point to the same Arc
+ let projection_exprs = vec![
+ ProjectionExpr {
+ expr: Arc::clone(&col_expr),
+ alias: "a1".to_string(),
+ },
+ ProjectionExpr {
+ expr: Arc::clone(&col_expr), // Same Arc!
+ alias: "a2".to_string(),
+ },
+ ];
+
+ let input = Arc::new(EmptyExec::new(schema));
+ let exec_plan = Arc::new(ProjectionExec::try_new(projection_exprs, input)?);
+
+ let ctx = SessionContext::new();
+ let codec = DefaultPhysicalExtensionCodec {};
+ let proto_converter = DeduplicatingProtoConverter {};
+
+ // Serialize
+ let bytes = physical_plan_to_bytes_with_proto_converter(
+ Arc::clone(&exec_plan) as Arc<dyn ExecutionPlan>,
+ &codec,
+ &proto_converter,
+ )?;
+
+ // Deserialize with a fresh converter
+ let deser_converter = DeduplicatingProtoConverter {};
+ let result_plan = physical_plan_from_bytes_with_proto_converter(
+ bytes.as_ref(),
+ ctx.task_ctx().as_ref(),
+ &codec,
+ &deser_converter,
+ )?;
+
+ // Get the projection from the result
+ let projection = result_plan
+ .as_any()
+ .downcast_ref::<ProjectionExec>()
+ .expect("Expected ProjectionExec");
+
+ let exprs: Vec<_> = projection.expr().iter().collect();
+ assert_eq!(exprs.len(), 2);
+
+ // The key test: both expressions should point to the same Arc after deduplication
+ // This is because they were the same Arc before serialization
+ assert!(
+ Arc::ptr_eq(&exprs[0].expr, &exprs[1].expr),
+ "Expected both expressions to share the same Arc after deduplication"
+ );
+
+ Ok(())
+}
+
+/// Test backward compatibility: protos without expr_id should still deserialize correctly.
+#[test]
+fn test_backward_compatibility_no_expr_id() -> Result<()> {
+ let field_a = Field::new("a", DataType::Int64, false);
+ let schema = Arc::new(Schema::new(vec![field_a]));
+
+ // Manually create a proto without expr_id set
+ let proto = PhysicalExprNode {
+ expr_id: None, // Simulating old proto without this field
+ expr_type: Some(
+ datafusion_proto::protobuf::physical_expr_node::ExprType::Column(
+ datafusion_proto::protobuf::PhysicalColumn {
+ name: "a".to_string(),
+ index: 0,
+ },
+ ),
+ ),
+ };
+
+ let ctx = SessionContext::new();
+ let codec = DefaultPhysicalExtensionCodec {};
+ let proto_converter = DefaultPhysicalProtoConverter {};
+
+ // Should deserialize without error
+ let result = proto_converter.proto_to_physical_expr(
+ &proto,
+ ctx.task_ctx().as_ref(),
+ &schema,
+ &codec,
+ )?;
+
+ // Verify the result is correct
+ let col = result
+ .as_any()
+ .downcast_ref::<Column>()
+ .expect("Expected Column");
+ assert_eq!(col.name(), "a");
+ assert_eq!(col.index(), 0);
+
+ Ok(())
+}
+
+/// Test that deduplication works within a single plan deserialization and that
+/// separate deserializations produce independent expressions (no cross-operation sharing).
+#[test]
+fn test_deduplication_within_plan_deserialization() -> Result<()> {
+ use datafusion_proto::bytes::{
+ physical_plan_from_bytes_with_proto_converter,
+ physical_plan_to_bytes_with_proto_converter,
+ };
+
+ let field_a = Field::new("a", DataType::Int64, false);
+ let schema = Arc::new(Schema::new(vec![field_a]));
+
+ // Create a plan with expressions that will be deduplicated
+ let col_expr: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
+ let projection_exprs = vec![
+ ProjectionExpr {
+ expr: Arc::clone(&col_expr),
+ alias: "a1".to_string(),
+ },
+ ProjectionExpr {
+ expr: Arc::clone(&col_expr), // Same Arc - will be deduplicated
+ alias: "a2".to_string(),
+ },
+ ];
+ let exec_plan = Arc::new(ProjectionExec::try_new(
+ projection_exprs,
+ Arc::new(EmptyExec::new(schema)),
+ )?);
+
+ let ctx = SessionContext::new();
+ let codec = DefaultPhysicalExtensionCodec {};
+ let proto_converter = DeduplicatingProtoConverter {};
+
+ // Serialize
+ let bytes = physical_plan_to_bytes_with_proto_converter(
+ Arc::clone(&exec_plan) as Arc<dyn ExecutionPlan>,
+ &codec,
+ &proto_converter,
+ )?;
+
+ // First deserialization
+ let plan1 = physical_plan_from_bytes_with_proto_converter(
+ bytes.as_ref(),
+ ctx.task_ctx().as_ref(),
+ &codec,
+ &proto_converter,
+ )?;
+
+ // Check that the plan was deserialized correctly with deduplication
+ let projection1 = plan1
+ .as_any()
+ .downcast_ref::<ProjectionExec>()
+ .expect("Expected ProjectionExec");
+ let exprs1: Vec<_> = projection1.expr().iter().collect();
+ assert_eq!(exprs1.len(), 2);
+ assert!(
+ Arc::ptr_eq(&exprs1[0].expr, &exprs1[1].expr),
+ "Expected both expressions to share the same Arc after deduplication"
+ );
+
+ // Second deserialization
+ let plan2 = physical_plan_from_bytes_with_proto_converter(
+ bytes.as_ref(),
+ ctx.task_ctx().as_ref(),
+ &codec,
+ &proto_converter,
+ )?;
+
+ // Check that the second plan was also deserialized correctly
+ let projection2 = plan2
+ .as_any()
+ .downcast_ref::<ProjectionExec>()
+ .expect("Expected ProjectionExec");
+ let exprs2: Vec<_> = projection2.expr().iter().collect();
+ assert_eq!(exprs2.len(), 2);
+ assert!(
+ Arc::ptr_eq(&exprs2[0].expr, &exprs2[1].expr),
+ "Expected both expressions to share the same Arc after deduplication"
+ );
+
+ // Check that there was no deduplication across deserializations
+ assert!(
+ !Arc::ptr_eq(&exprs1[0].expr, &exprs2[0].expr),
+ "Expected expressions from different deserializations to be different Arcs"
+ );
+ assert!(
+ !Arc::ptr_eq(&exprs1[1].expr, &exprs2[1].expr),
+ "Expected expressions from different deserializations to be different Arcs"
+ );
+
+ Ok(())
+}
+
+/// Test that deduplication works within direct expression deserialization and that
+/// separate deserializations produce independent expressions (no cross-operation sharing).
+#[test]
+fn test_deduplication_within_expr_deserialization() -> Result<()> {
+ let field_a = Field::new("a", DataType::Int64, false);
+ let schema = Arc::new(Schema::new(vec![field_a]));
+
+ // Create a binary expression where both sides are the same Arc
+ // This allows us to test deduplication within a single deserialization
+ let col_expr: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
+ let binary_expr: Arc<dyn PhysicalExpr> = Arc::new(BinaryExpr::new(
+ Arc::clone(&col_expr),
+ Operator::Plus,
+ Arc::clone(&col_expr), // Same Arc - will be deduplicated
+ ));
+
+ let ctx = SessionContext::new();
+ let codec = DefaultPhysicalExtensionCodec {};
+ let proto_converter = DeduplicatingProtoConverter {};
+
+ // Serialize the expression
+ let proto = proto_converter.physical_expr_to_proto(&binary_expr, &codec)?;
+
+ // First expression deserialization
+ let expr1 = proto_converter.proto_to_physical_expr(
+ &proto,
+ ctx.task_ctx().as_ref(),
+ &schema,
+ &codec,
+ )?;
+
+ // Check that deduplication worked within the deserialization
+ let binary1 = expr1
+ .as_any()
+ .downcast_ref::<BinaryExpr>()
+ .expect("Expected BinaryExpr");
+ assert!(
+ Arc::ptr_eq(binary1.left(), binary1.right()),
+ "Expected both sides to share the same Arc after deduplication"
+ );
+
+ // Second expression deserialization
+ let expr2 = proto_converter.proto_to_physical_expr(
+ &proto,
+ ctx.task_ctx().as_ref(),
+ &schema,
+ &codec,
+ )?;
+
+ // Check that the second expression was also deserialized correctly
+ let binary2 = expr2
+ .as_any()
+ .downcast_ref::<BinaryExpr>()
+ .expect("Expected BinaryExpr");
+ assert!(
+ Arc::ptr_eq(binary2.left(), binary2.right()),
+ "Expected both sides to share the same Arc after deduplication"
+ );
+
+ // Check that there was no deduplication across deserializations
+ assert!(
+ !Arc::ptr_eq(binary1.left(), binary2.left()),
+ "Expected expressions from different deserializations to be different Arcs"
+ );
+ assert!(
+ !Arc::ptr_eq(binary1.right(), binary2.right()),
+ "Expected expressions from different deserializations to be different Arcs"
+ );
+
+ Ok(())
+}
+
+/// Test that session_id rotates between top-level serialization operations.
+/// This verifies that each top-level serialization gets a fresh session_id,
+/// which prevents cross-process collisions when serialized plans are merged.
+#[test]
+fn test_session_id_rotation_between_serializations() -> Result<()> {
+ let field_a = Field::new("a", DataType::Int64, false);
+ let _schema = Arc::new(Schema::new(vec![field_a]));
+
+ // Create a simple expression
+ let col_expr: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
+
+ let codec = DefaultPhysicalExtensionCodec {};
+ let proto_converter = DeduplicatingProtoConverter {};
+
+ // First serialization
+ let proto1 = proto_converter.physical_expr_to_proto(&col_expr, &codec)?;
+ let expr_id1 = proto1.expr_id.expect("Expected expr_id to be set");
+
+ // Second serialization with the same converter
+ // The session_id should have rotated, so the expr_id should be different
+ // even though we're serializing the same expression (same pointer address)
+ let proto2 = proto_converter.physical_expr_to_proto(&col_expr, &codec)?;
+ let expr_id2 = proto2.expr_id.expect("Expected expr_id to be set");
+
+ // The expr_ids should be different because session_id rotated
+ assert_ne!(
+ expr_id1, expr_id2,
+ "Expected different expr_ids due to session_id rotation between serializations"
+ );
+
+ // Also test that serializing the same expression multiple times within
+ // the same top-level operation would give the same expr_id (not testable
+ // here directly since each physical_expr_to_proto is a top-level operation,
+ // but the deduplication tests verify this indirectly)
+
+ Ok(())
+}
+
+/// Test that session_id rotation works correctly with execution plans.
+/// This verifies the end-to-end behavior with plan serialization.
+#[test]
+fn test_session_id_rotation_with_execution_plans() -> Result<()> {
+ use datafusion_proto::bytes::physical_plan_to_bytes_with_proto_converter;
+
+ let field_a = Field::new("a", DataType::Int64, false);
+ let schema = Arc::new(Schema::new(vec![field_a]));
+
+ // Create a simple plan
+ let col_expr: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
+ let projection_exprs = vec![ProjectionExpr {
+ expr: Arc::clone(&col_expr),
+ alias: "a1".to_string(),
+ }];
+ let exec_plan = Arc::new(ProjectionExec::try_new(
+ projection_exprs.clone(),
+ Arc::new(EmptyExec::new(Arc::clone(&schema))),
+ )?);
+
+ let codec = DefaultPhysicalExtensionCodec {};
+ let proto_converter = DeduplicatingProtoConverter {};
+
+ // First serialization
+ let bytes1 = physical_plan_to_bytes_with_proto_converter(
+ Arc::clone(&exec_plan) as Arc<dyn ExecutionPlan>,
+ &codec,
+ &proto_converter,
+ )?;
+
+ // Second serialization with the same converter
+ let bytes2 = physical_plan_to_bytes_with_proto_converter(
+ Arc::clone(&exec_plan) as Arc<dyn ExecutionPlan>,
+ &codec,
+ &proto_converter,
+ )?;
+
+ // The serialized bytes should be different due to different session_ids
+ // (specifically, the expr_id values embedded in the protobuf will differ)
+ assert_ne!(
+ bytes1.as_ref(),
+ bytes2.as_ref(),
+ "Expected different serialized bytes due to session_id rotation"
+ );
+
+ // But both should deserialize correctly
+ let ctx = SessionContext::new();
+ let deser_converter = DeduplicatingProtoConverter {};
+
+ let plan1 = datafusion_proto::bytes::physical_plan_from_bytes_with_proto_converter(
+ bytes1.as_ref(),
+ ctx.task_ctx().as_ref(),
+ &codec,
+ &deser_converter,
+ )?;
+
+ let plan2 = datafusion_proto::bytes::physical_plan_from_bytes_with_proto_converter(
+ bytes2.as_ref(),
+ ctx.task_ctx().as_ref(),
+ &codec,
+ &deser_converter,
+ )?;
+
+ // Verify both plans have the expected structure
+ assert_eq!(plan1.schema(), plan2.schema());
+
+ Ok(())
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org