diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs index a91903934ac9..fa72e993637f 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs @@ -677,6 +677,114 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { fn metrics(&self) -> Option { Some(self.metrics.clone_inner()) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let buffered = ctx.encode_child(self.buffered())?; + let streamed = ctx.encode_child(self.streamed())?; + let on_buffered = ctx.encode_expr(&self.on.0)?; + let on_streamed = ctx.encode_expr(&self.on.1)?; + let join_type = crate::joins::proto::join_type_to_proto(self.join_type()); + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::PiecewiseMergeJoin( + Box::new(protobuf::PiecewiseMergeJoinExecNode { + buffered: Some(Box::new(buffered)), + streamed: Some(Box::new(streamed)), + on_buffered: Some(on_buffered), + on_streamed: Some(on_streamed), + // Matches the `Operator` encoding used for `BinaryExpr`: + // the `Debug` name of the variant. + operator: format!("{:?}", self.operator), + join_type: join_type.into(), + num_partitions: self.num_partitions as u64, + }), + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl PiecewiseMergeJoinExec { + /// Reconstruct a [`PiecewiseMergeJoinExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]. Every other field of + /// the operator (schema, sort options, required orderings, plan properties) is + /// derived by [`PiecewiseMergeJoinExec::try_new`], so it is not on the wire. + /// + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_common::{internal_datafusion_err, plan_datafusion_err}; + use datafusion_proto_models::protobuf; + + let join = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::PiecewiseMergeJoin, + "PiecewiseMergeJoinExec", + ); + let buffered = ctx.decode_required_child( + join.buffered.as_deref(), + "PiecewiseMergeJoinExec", + "buffered", + )?; + let streamed = ctx.decode_required_child( + join.streamed.as_deref(), + "PiecewiseMergeJoinExec", + "streamed", + )?; + let on_buffered = ctx.decode_required_expr( + join.on_buffered.as_ref(), + buffered.schema().as_ref(), + "PiecewiseMergeJoinExec", + "on_buffered", + )?; + let on_streamed = ctx.decode_required_expr( + join.on_streamed.as_ref(), + streamed.schema().as_ref(), + "PiecewiseMergeJoinExec", + "on_streamed", + )?; + + let operator = Operator::from_proto_name(&join.operator).ok_or_else(|| { + internal_datafusion_err!( + "PiecewiseMergeJoinExec: unknown Operator '{}'", + join.operator + ) + })?; + let join_type = crate::joins::proto::join_type_from_proto( + join.join_type, + "PiecewiseMergeJoinExec", + )?; + + // Checked rather than `as usize`: a truncated partition count would not + // fail loudly, it would silently change how the buffered side is split. + let num_partitions = + usize::try_from(join.num_partitions).map_err(|_| { + plan_datafusion_err!( + "PiecewiseMergeJoinExec: num_partitions {} cannot be represented as usize on this target", + join.num_partitions + ) + })?; + + Ok(Arc::new(Self::try_new( + buffered, + streamed, + (on_buffered, on_streamed), + operator, + join_type, + num_partitions, + )?)) + } } impl DisplayAs for PiecewiseMergeJoinExec { diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 43a90264c2b1..784ec4e0e1d1 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -899,6 +899,7 @@ message PhysicalPlanNode { BufferExecNode buffer = 37; ArrowScanExecNode arrow_scan = 38; ScalarSubqueryExecNode scalar_subquery = 39; + PiecewiseMergeJoinExecNode piecewise_merge_join = 40; } } @@ -1678,6 +1679,18 @@ message SortMergeJoinExecNode { datafusion_common.NullEquality null_equality = 7; } +message PiecewiseMergeJoinExecNode { + PhysicalPlanNode buffered = 1; + PhysicalPlanNode streamed = 2; + // The buffered-side and streamed-side halves of the single range predicate. + PhysicalExprNode on_buffered = 3; + PhysicalExprNode on_streamed = 4; + // `Operator` variant name, e.g. "Lt". Must be one of Lt/LtEq/Gt/GtEq. + string operator = 5; + datafusion_common.JoinType join_type = 6; + uint64 num_partitions = 7; +} + message AsyncFuncExecNode { PhysicalPlanNode input = 1; repeated PhysicalExprNode async_exprs = 2; diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 908f9752b7f1..9b4f39991785 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -20419,6 +20419,9 @@ impl serde::Serialize for PhysicalPlanNode { physical_plan_node::PhysicalPlanType::ScalarSubquery(v) => { struct_ser.serialize_field("scalarSubquery", v)?; } + physical_plan_node::PhysicalPlanType::PiecewiseMergeJoin(v) => { + struct_ser.serialize_field("piecewiseMergeJoin", v)?; + } } } struct_ser.end() @@ -20491,6 +20494,8 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { "arrowScan", "scalar_subquery", "scalarSubquery", + "piecewise_merge_join", + "piecewiseMergeJoin", ]; #[allow(clippy::enum_variant_names)] @@ -20533,6 +20538,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { Buffer, ArrowScan, ScalarSubquery, + PiecewiseMergeJoin, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -20592,6 +20598,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { "buffer" => Ok(GeneratedField::Buffer), "arrowScan" | "arrow_scan" => Ok(GeneratedField::ArrowScan), "scalarSubquery" | "scalar_subquery" => Ok(GeneratedField::ScalarSubquery), + "piecewiseMergeJoin" | "piecewise_merge_join" => Ok(GeneratedField::PiecewiseMergeJoin), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -20878,6 +20885,13 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { return Err(serde::de::Error::duplicate_field("scalarSubquery")); } physical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_plan_node::PhysicalPlanType::ScalarSubquery) +; + } + GeneratedField::PiecewiseMergeJoin => { + if physical_plan_type__.is_some() { + return Err(serde::de::Error::duplicate_field("piecewiseMergeJoin")); + } + physical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_plan_node::PhysicalPlanType::PiecewiseMergeJoin) ; } } @@ -22203,6 +22217,209 @@ impl<'de> serde::Deserialize<'de> for PhysicalWindowExprNode { deserializer.deserialize_struct("datafusion.PhysicalWindowExprNode", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for PiecewiseMergeJoinExecNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.buffered.is_some() { + len += 1; + } + if self.streamed.is_some() { + len += 1; + } + if self.on_buffered.is_some() { + len += 1; + } + if self.on_streamed.is_some() { + len += 1; + } + if !self.operator.is_empty() { + len += 1; + } + if self.join_type != 0 { + len += 1; + } + if self.num_partitions != 0 { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.PiecewiseMergeJoinExecNode", len)?; + if let Some(v) = self.buffered.as_ref() { + struct_ser.serialize_field("buffered", v)?; + } + if let Some(v) = self.streamed.as_ref() { + struct_ser.serialize_field("streamed", v)?; + } + if let Some(v) = self.on_buffered.as_ref() { + struct_ser.serialize_field("onBuffered", v)?; + } + if let Some(v) = self.on_streamed.as_ref() { + struct_ser.serialize_field("onStreamed", v)?; + } + if !self.operator.is_empty() { + struct_ser.serialize_field("operator", &self.operator)?; + } + if self.join_type != 0 { + let v = super::datafusion_common::JoinType::try_from(self.join_type) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.join_type)))?; + struct_ser.serialize_field("joinType", &v)?; + } + if self.num_partitions != 0 { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("numPartitions", ToString::to_string(&self.num_partitions).as_str())?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for PiecewiseMergeJoinExecNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "buffered", + "streamed", + "on_buffered", + "onBuffered", + "on_streamed", + "onStreamed", + "operator", + "join_type", + "joinType", + "num_partitions", + "numPartitions", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Buffered, + Streamed, + OnBuffered, + OnStreamed, + Operator, + JoinType, + NumPartitions, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "buffered" => Ok(GeneratedField::Buffered), + "streamed" => Ok(GeneratedField::Streamed), + "onBuffered" | "on_buffered" => Ok(GeneratedField::OnBuffered), + "onStreamed" | "on_streamed" => Ok(GeneratedField::OnStreamed), + "operator" => Ok(GeneratedField::Operator), + "joinType" | "join_type" => Ok(GeneratedField::JoinType), + "numPartitions" | "num_partitions" => Ok(GeneratedField::NumPartitions), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = PiecewiseMergeJoinExecNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.PiecewiseMergeJoinExecNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut buffered__ = None; + let mut streamed__ = None; + let mut on_buffered__ = None; + let mut on_streamed__ = None; + let mut operator__ = None; + let mut join_type__ = None; + let mut num_partitions__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Buffered => { + if buffered__.is_some() { + return Err(serde::de::Error::duplicate_field("buffered")); + } + buffered__ = map_.next_value()?; + } + GeneratedField::Streamed => { + if streamed__.is_some() { + return Err(serde::de::Error::duplicate_field("streamed")); + } + streamed__ = map_.next_value()?; + } + GeneratedField::OnBuffered => { + if on_buffered__.is_some() { + return Err(serde::de::Error::duplicate_field("onBuffered")); + } + on_buffered__ = map_.next_value()?; + } + GeneratedField::OnStreamed => { + if on_streamed__.is_some() { + return Err(serde::de::Error::duplicate_field("onStreamed")); + } + on_streamed__ = map_.next_value()?; + } + GeneratedField::Operator => { + if operator__.is_some() { + return Err(serde::de::Error::duplicate_field("operator")); + } + operator__ = Some(map_.next_value()?); + } + GeneratedField::JoinType => { + if join_type__.is_some() { + return Err(serde::de::Error::duplicate_field("joinType")); + } + join_type__ = Some(map_.next_value::()? as i32); + } + GeneratedField::NumPartitions => { + if num_partitions__.is_some() { + return Err(serde::de::Error::duplicate_field("numPartitions")); + } + num_partitions__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } + } + } + Ok(PiecewiseMergeJoinExecNode { + buffered: buffered__, + streamed: streamed__, + on_buffered: on_buffered__, + on_streamed: on_streamed__, + operator: operator__.unwrap_or_default(), + join_type: join_type__.unwrap_or_default(), + num_partitions: num_partitions__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.PiecewiseMergeJoinExecNode", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for PlaceholderNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index ba00577ab9a1..2f8a53010327 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -1342,7 +1342,7 @@ pub mod table_reference { pub struct PhysicalPlanNode { #[prost( oneof = "physical_plan_node::PhysicalPlanType", - tags = "1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39" + tags = "1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40" )] pub physical_plan_type: ::core::option::Option, } @@ -1428,6 +1428,10 @@ pub mod physical_plan_node { ArrowScan(super::ArrowScanExecNode), #[prost(message, tag = "39")] ScalarSubquery(::prost::alloc::boxed::Box), + #[prost(message, tag = "40")] + PiecewiseMergeJoin( + ::prost::alloc::boxed::Box, + ), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -2550,6 +2554,25 @@ pub struct SortMergeJoinExecNode { pub null_equality: i32, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct PiecewiseMergeJoinExecNode { + #[prost(message, optional, boxed, tag = "1")] + pub buffered: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, optional, boxed, tag = "2")] + pub streamed: ::core::option::Option<::prost::alloc::boxed::Box>, + /// The buffered-side and streamed-side halves of the single range predicate. + #[prost(message, optional, tag = "3")] + pub on_buffered: ::core::option::Option, + #[prost(message, optional, tag = "4")] + pub on_streamed: ::core::option::Option, + /// `Operator` variant name, e.g. "Lt". Must be one of Lt/LtEq/Gt/GtEq. + #[prost(string, tag = "5")] + pub operator: ::prost::alloc::string::String, + #[prost(enumeration = "super::datafusion_common::JoinType", tag = "6")] + pub join_type: i32, + #[prost(uint64, tag = "7")] + pub num_partitions: u64, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct AsyncFuncExecNode { #[prost(message, optional, boxed, tag = "1")] pub input: ::core::option::Option<::prost::alloc::boxed::Box>, diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 222901aff521..6865f032dea6 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -60,8 +60,8 @@ use datafusion_physical_plan::empty::EmptyExec; use datafusion_physical_plan::explain::ExplainExec; use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::{ - CrossJoinExec, HashJoinExec, NestedLoopJoinExec, SortMergeJoinExec, - SymmetricHashJoinExec, + CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PiecewiseMergeJoinExec, + SortMergeJoinExec, SymmetricHashJoinExec, }; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::memory::LazyMemoryExec; @@ -1209,6 +1209,9 @@ pub trait PhysicalPlanNodeExt: Sized { PhysicalPlanType::ScalarSubquery(_) => { ScalarSubqueryExec::try_from_proto(self.node(), &decode_ctx) } + PhysicalPlanType::PiecewiseMergeJoin(_) => { + PiecewiseMergeJoinExec::try_from_proto(self.node(), &decode_ctx) + } } } diff --git a/datafusion/proto/tests/cases/plans/joins.rs b/datafusion/proto/tests/cases/plans/joins.rs index 941e8832952c..24e2a32560d1 100644 --- a/datafusion/proto/tests/cases/plans/joins.rs +++ b/datafusion/proto/tests/cases/plans/joins.rs @@ -24,17 +24,29 @@ use datafusion::logical_expr::{JoinType, Operator}; use datafusion::physical_expr::LexOrdering; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::empty::EmptyExec; -use datafusion::physical_plan::expressions::{BinaryExpr, Column, PhysicalSortExpr}; +use datafusion::physical_plan::expressions::{ + BinaryExpr, CastExpr, Column, Literal, PhysicalSortExpr, +}; use datafusion::physical_plan::joins::utils::{ColumnIndex, JoinFilter}; use datafusion::physical_plan::joins::{ - HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, - StreamJoinPartitionMode, SymmetricHashJoinExec, + HashJoinExec, NestedLoopJoinExec, PartitionMode, PiecewiseMergeJoinExec, + SortMergeJoinExec, StreamJoinPartitionMode, SymmetricHashJoinExec, }; use datafusion::prelude::SessionContext; +use datafusion_common::ScalarValue; +use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_common::{JoinSide, NullEquality, Result}; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_proto::bytes::{ + physical_plan_from_bytes_with_proto_converter, + physical_plan_to_bytes_with_proto_converter, +}; use datafusion_proto::physical_plan::{ DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, }; +use datafusion_proto::protobuf; +use datafusion_proto::protobuf::physical_plan_node::PhysicalPlanType; +use prost::Message; use std::sync::Arc; use std::vec; @@ -456,3 +468,535 @@ async fn roundtrip_logical_plan_sort_merge_join() -> Result<()> { let plan = ctx.sql(query).await?.create_physical_plan().await?; roundtrip_test(plan) } + +/// One `roundtrip_piecewise_merge_join_compound_on_exprs` case: the buffered and +/// streamed halves of the predicate, and the string each must display as after the +/// round trip. +type PiecewiseOnCase = ( + Arc, + Arc, + &'static str, + &'static str, +); + +/// A field of `PiecewiseMergeJoinExecNode` to clear, with its wire name. +type PiecewiseFieldClear = (&'static str, fn(&mut protobuf::PiecewiseMergeJoinExecNode)); + +/// Schemas for the `PiecewiseMergeJoinExec` tests. +/// +/// The two sides differ in width, in column names, and in the *index* of the join +/// column (buffered `a` is at 2, streamed `b` is at 0). Matching indices on both +/// sides would let a decoder that resolved a predicate half against the wrong +/// schema still produce `Column { index: 0 }` and round-trip silently. +fn piecewise_schemas() -> (Arc, Arc) { + ( + Arc::new(Schema::new(vec![ + Field::new("pad0", DataType::Utf8, true), + Field::new("pad1", DataType::Int64, false), + Field::new("a", DataType::Int64, false), + ])), + Arc::new(Schema::new(vec![ + Field::new("b", DataType::Int64, true), + Field::new("pad2", DataType::Utf8, false), + ])), + ) +} + +/// A valid `PiecewiseMergeJoinExec` over [`piecewise_schemas`], for tests that then +/// corrupt one field of its encoding. +fn piecewise_join( + schemas: &(Arc, Arc), + num_partitions: usize, +) -> Result> { + Ok(Arc::new(PiecewiseMergeJoinExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schemas.0))), + Arc::new(EmptyExec::new(Arc::clone(&schemas.1))), + ( + Arc::new(Column::new("a", 2)) as _, + Arc::new(Column::new("b", 0)) as _, + ), + Operator::Lt, + JoinType::Inner, + num_partitions, + )?)) +} + +/// `PiecewiseMergeJoinExec` derives its schema, sort options, required input +/// orderings and plan properties inside `try_new`, so only the six constructor +/// arguments travel on the wire. Cover the full cartesian product of the two +/// enum-valued ones: every range operator against every supported join type -- +/// the four classic ones plus the two left existence joins, whose output schema is +/// the buffered side alone. (Right existence joins and Mark joins are rejected by +/// `try_new`, so this is the complete set.) +#[test] +fn roundtrip_piecewise_merge_join() -> Result<()> { + let (schema_buffered, schema_streamed) = piecewise_schemas(); + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + + for operator in [Operator::Lt, Operator::LtEq, Operator::Gt, Operator::GtEq] { + for join_type in [ + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::Full, + JoinType::LeftSemi, + JoinType::LeftAnti, + ] { + let result = roundtrip_test_and_return( + Arc::new(PiecewiseMergeJoinExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema_buffered))), + Arc::new(EmptyExec::new(Arc::clone(&schema_streamed))), + ( + Arc::new(Column::new("a", 2)) as _, + Arc::new(Column::new("b", 0)) as _, + ), + operator, + join_type, + 7, + )?), + &ctx, + &codec, + &proto_converter, + )?; + let result = result.downcast_ref::().unwrap(); + + assert_eq!(result.operator, operator); + assert_eq!(result.join_type(), join_type); + // Both the name and the index have to survive on the correct side. + assert_eq!(result.on.0.to_string(), "a@2"); + assert_eq!(result.on.1.to_string(), "b@0"); + // The existence joins output the buffered side alone (3 fields) while the + // classic ones output both sides (3 + 2). The before/after comparison + // inside the helper already covers the schema; this pins the expected + // width absolutely, so the existence output contract is stated rather + // than merely preserved. + let expected_fields = + if matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti) { + 3 + } else { + 5 + }; + assert_eq!( + result.schema().fields().len(), + expected_fields, + "unexpected output width for {join_type}" + ); + } + } + Ok(()) +} + +/// The two halves of the range predicate are arbitrary `PhysicalExpr`s, not just +/// columns: `side_of` in the physical planner classifies a side by +/// `Expr::column_refs()`, so any expression whose columns all come from one input +/// qualifies and reaches `create_physical_expr`. `ON t0.a + 1 < t1.b * 2` therefore +/// puts a `BinaryExpr` tree on each side. Cover a nested tree and a `CastExpr` so +/// the `encode_expr` / `decode_required_expr` path is exercised beyond a bare +/// `Column`, including the literals inside it. +#[test] +fn roundtrip_piecewise_merge_join_compound_on_exprs() -> Result<()> { + let (schema_buffered, schema_streamed) = piecewise_schemas(); + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + + // (buffered half, streamed half, expected buffered string, expected streamed string) + let cases: Vec = vec![ + // Nested arithmetic on both sides. + ( + Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 2)), + Operator::Plus, + Arc::new(Literal::new(ScalarValue::Int64(Some(1)))), + )), + Operator::Multiply, + Arc::new(Column::new("pad1", 1)), + )), + Arc::new(BinaryExpr::new( + Arc::new(Column::new("b", 0)), + Operator::Minus, + Arc::new(Literal::new(ScalarValue::Int64(Some(2)))), + )), + "(a@2 + 1) * pad1@1", + "b@0 - 2", + ), + // A cast on the buffered side, a plain column on the streamed side: the + // two halves need not have the same shape. + ( + Arc::new(CastExpr::new( + Arc::new(Column::new("a", 2)), + DataType::Int32, + None, + )), + Arc::new(Column::new("b", 0)), + "CAST(a@2 AS Int32)", + "b@0", + ), + ]; + + for (on_buffered, on_streamed, expected_buffered, expected_streamed) in cases { + let result = roundtrip_test_and_return( + Arc::new(PiecewiseMergeJoinExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema_buffered))), + Arc::new(EmptyExec::new(Arc::clone(&schema_streamed))), + (Arc::clone(&on_buffered), Arc::clone(&on_streamed)), + Operator::Lt, + JoinType::Inner, + 7, + )?), + &ctx, + &codec, + &proto_converter, + )?; + let result = result.downcast_ref::().unwrap(); + + assert_eq!(result.on.0.to_string(), expected_buffered); + assert_eq!(result.on.1.to_string(), expected_streamed); + } + Ok(()) +} + +/// `num_partitions` crosses the wire as a `u64` and is read back with a checked +/// `usize::try_from`, matching how `HashJoinExec` handles `fetch`. Pin the +/// conversion at the boundaries rather than only at a small value. +/// +/// Both `u32::MAX` and `usize::MAX` are representable on every target (on a 32-bit +/// target `usize::MAX` is simply `u32::MAX`), so this stays portable. The +/// truncating case -- a `u64` above `usize::MAX` -- is only reachable on a 32-bit +/// target and so is not exercised on a 64-bit host. `num_partitions` is only read +/// at execution time, so constructing these does not allocate. +#[test] +fn roundtrip_piecewise_merge_join_num_partitions_bounds() -> Result<()> { + let schemas = piecewise_schemas(); + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + + for num_partitions in [1, 7, u32::MAX as usize, usize::MAX] { + // `roundtrip_test_and_return` compares the `Debug` output, which for + // `PiecewiseMergeJoinExec` includes `num_partitions`. + roundtrip_test_and_return( + piecewise_join(&schemas, num_partitions)?, + &ctx, + &codec, + &proto_converter, + )?; + } + Ok(()) +} + +/// Every message field except `operator` and `join_type` is a proto3 `message`, +/// so "absent" is representable on the wire and a truncated or hand-built payload +/// can omit any of them. Each omission must name the field it is missing rather +/// than panicking on an `unwrap`. +#[test] +fn piecewise_merge_join_rejects_missing_fields() -> Result<()> { + let schemas = piecewise_schemas(); + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + + let valid = physical_plan_to_bytes_with_proto_converter( + piecewise_join(&schemas, 7)?, + &codec, + &proto_converter, + )?; + + let clear: [PiecewiseFieldClear; 4] = [ + ("buffered", |join| join.buffered = None), + ("streamed", |join| join.streamed = None), + ("on_buffered", |join| join.on_buffered = None), + ("on_streamed", |join| join.on_streamed = None), + ]; + + for (field, clear_field) in clear { + let mut node = protobuf::PhysicalPlanNode::decode(valid.as_ref()) + .expect("a plan encoded by try_to_proto must decode as a PhysicalPlanNode"); + let Some(PhysicalPlanType::PiecewiseMergeJoin(join)) = + node.physical_plan_type.as_mut() + else { + panic!("expected a PiecewiseMergeJoin node"); + }; + clear_field(join); + + let Err(err) = physical_plan_from_bytes_with_proto_converter( + &node.encode_to_vec(), + ctx.task_ctx().as_ref(), + &codec, + &proto_converter, + ) else { + panic!("decoding must fail when {field} is absent"); + }; + // Match the quoted field name, so that clearing `buffered` cannot be + // satisfied by an error that only mentions `on_buffered`. + let expected = format!("missing required field '{field}'"); + assert!( + err.to_string().contains(&expected), + "missing {field}: expected an error containing {expected:?}, got: {err}" + ); + } + Ok(()) +} + +/// The operator travels as its `Operator` variant name, so the decoder has to +/// handle names that `try_to_proto` would never emit: a name with no `Operator` +/// counterpart, and a real `Operator` that is not one of the four range +/// operators. Both must surface as errors rather than panicking or silently +/// decoding into a different operator, since a malformed payload can arrive from +/// any peer. +#[test] +fn piecewise_merge_join_rejects_bad_operator_on_the_wire() -> Result<()> { + let schemas = piecewise_schemas(); + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + + // Start from a valid encoding so only the operator field is malformed. + let valid = physical_plan_to_bytes_with_proto_converter( + piecewise_join(&schemas, 7)?, + &codec, + &proto_converter, + )?; + let mut node = protobuf::PhysicalPlanNode::decode(valid.as_ref()) + .expect("a plan just encoded by try_to_proto must decode as a PhysicalPlanNode"); + + for (operator, expected) in [ + ("NotAnOperator", "unknown Operator"), + ("Eq", "non-range operator"), + ] { + let Some(PhysicalPlanType::PiecewiseMergeJoin(join)) = + node.physical_plan_type.as_mut() + else { + panic!("expected a PiecewiseMergeJoin node"); + }; + join.operator = operator.to_string(); + + let Err(err) = physical_plan_from_bytes_with_proto_converter( + &node.encode_to_vec(), + ctx.task_ctx().as_ref(), + &codec, + &proto_converter, + ) else { + panic!("decoding must fail for operator {operator}"); + }; + assert!( + err.to_string().contains(expected), + "operator {operator}: expected an error containing {expected:?}, got: {err}" + ); + } + Ok(()) +} + +/// `join_type` is a proto3 enum, so any `i32` is representable on the wire -- including +/// the existence joins `try_new` still rejects (right-sided and mark) and values that +/// map to no `JoinType` at all. `try_to_proto` can emit none of these, but a payload +/// from a newer peer or a hand-built one can, and each must surface as an error +/// rather than a panic or a silently different operator. The right-sided cases matter +/// most: `try_new` derives *reversed* sort options for them, so accepting one would +/// build a plan whose buffered side is sorted the wrong way. +#[test] +fn piecewise_merge_join_rejects_unsupported_join_type_on_the_wire() -> Result<()> { + let schemas = piecewise_schemas(); + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + + let valid = physical_plan_to_bytes_with_proto_converter( + piecewise_join(&schemas, 7)?, + &codec, + &proto_converter, + )?; + let mut node = protobuf::PhysicalPlanNode::decode(valid.as_ref()) + .expect("a plan just encoded by try_to_proto must decode as a PhysicalPlanNode"); + + // (wire value, description, expected error fragment) + let cases: [(i32, &str, &str); 5] = [ + ( + protobuf::JoinType::Rightsemi as i32, + "RightSemi", + "Existence join RightSemi is currently not supported", + ), + ( + protobuf::JoinType::Rightanti as i32, + "RightAnti", + "Existence join RightAnti is currently not supported", + ), + ( + protobuf::JoinType::Leftmark as i32, + "LeftMark", + "Existence join LeftMark is currently not supported", + ), + ( + protobuf::JoinType::Rightmark as i32, + "RightMark", + "Existence join RightMark is currently not supported", + ), + // Past the last tag, so it maps to no variant. + (i32::MAX, "no variant", "unknown JoinType"), + ]; + + for (wire_value, description, expected) in cases { + let Some(PhysicalPlanType::PiecewiseMergeJoin(join)) = + node.physical_plan_type.as_mut() + else { + panic!("expected a PiecewiseMergeJoin node"); + }; + join.join_type = wire_value; + + let Err(err) = physical_plan_from_bytes_with_proto_converter( + &node.encode_to_vec(), + ctx.task_ctx().as_ref(), + &codec, + &proto_converter, + ) else { + panic!("decoding must fail for join_type {description}"); + }; + assert!( + err.to_string().contains(expected), + "join_type {description}: expected an error containing {expected:?}, got: {err}" + ); + } + Ok(()) +} + +/// Every other test here compares a plan against itself after a round trip, which is +/// blind to a *symmetric* mistake: if `try_to_proto` wrote the streamed side into the +/// `buffered` field and `try_from_proto` read it back the same way, before and after +/// would still match, and even the `on` assertions would hold, because each half is +/// resolved against whichever child the decoder paired it with. The bytes would still +/// be wrong for every other reader of this schema. So assert the wire layout +/// absolutely, once. A left existence join is the clearest case to do it on: its two +/// sides differ in width on the wire (3 fields buffered, 2 streamed), and its output +/// is the buffered side alone, so nothing downstream would reveal a swap either. +#[test] +fn piecewise_merge_join_existence_wire_layout() -> Result<()> { + let (schema_buffered, schema_streamed) = piecewise_schemas(); + + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + + let bytes = physical_plan_to_bytes_with_proto_converter( + Arc::new(PiecewiseMergeJoinExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema_buffered))), + Arc::new(EmptyExec::new(Arc::clone(&schema_streamed))), + ( + Arc::new(Column::new("a", 2)) as _, + Arc::new(Column::new("b", 0)) as _, + ), + Operator::Gt, + JoinType::LeftSemi, + 7, + )?), + &codec, + &proto_converter, + )?; + + let node = protobuf::PhysicalPlanNode::decode(bytes.as_ref()) + .expect("a plan encoded by try_to_proto must decode as a PhysicalPlanNode"); + let Some(PhysicalPlanType::PiecewiseMergeJoin(join)) = node.physical_plan_type else { + panic!("expected a PiecewiseMergeJoin node"); + }; + + // The two children, identified by the width of the schema each carries. + let child_width = |child: Option<&protobuf::PhysicalPlanNode>, field: &str| { + let Some(PhysicalPlanType::Empty(empty)) = child + .expect("child must be present") + .physical_plan_type + .as_ref() + else { + panic!("expected an EmptyExec under {field}"); + }; + empty + .schema + .as_ref() + .expect("schema must be present") + .columns + .len() + }; + assert_eq!(child_width(join.buffered.as_deref(), "buffered"), 3); + assert_eq!(child_width(join.streamed.as_deref(), "streamed"), 2); + + // Each half of the range predicate, against the side it belongs to. + let column = |expr: Option<&protobuf::PhysicalExprNode>, field: &str| { + let Some(protobuf::physical_expr_node::ExprType::Column(column)) = + expr.expect("expr must be present").expr_type.as_ref() + else { + panic!("expected a Column under {field}"); + }; + (column.name.clone(), column.index) + }; + assert_eq!( + column(join.on_buffered.as_ref(), "on_buffered"), + ("a".to_string(), 2) + ); + assert_eq!( + column(join.on_streamed.as_ref(), "on_streamed"), + ("b".to_string(), 0) + ); + + assert_eq!(join.operator, "Gt"); + assert_eq!(join.join_type, protobuf::JoinType::Leftsemi as i32); + assert_eq!(join.num_partitions, 7); + Ok(()) +} + +/// End-to-end: a `PiecewiseMergeJoinExec` as the planner actually builds it, rather +/// than one hand-constructed by the test. This covers what the unit test above +/// cannot: the planner may swap the two inputs and reverse the operator (a range +/// predicate written right-to-left becomes a left-to-right one on a swapped plan), +/// and it takes `num_partitions` from `target_partitions`. If the encoding confused +/// the buffered and streamed sides, the swapped form is where it would show. +#[tokio::test] +async fn roundtrip_planned_piecewise_merge_join() -> Result<()> { + let ctx = SessionContext::new(); + ctx.register_csv( + "t0", + "tests/testdata/test.csv", + datafusion::prelude::CsvReadOptions::default().has_header(true), + ) + .await?; + ctx.register_csv( + "t1", + "tests/testdata/test.csv", + datafusion::prelude::CsvReadOptions::default().has_header(true), + ) + .await?; + + ctx.sql("SET datafusion.optimizer.enable_piecewise_merge_join = true") + .await? + .collect() + .await?; + + // Both predicate directions, so the input-swapping branch of the planner is + // covered as well as the straight one, plus a compound predicate so the + // planner emits a non-`Column` expression on each side. The last two + // decorrelate to `LeftSemi` / `LeftAnti`, where the marked side must come back + // as the buffered one. + for query in [ + "SELECT t0.a FROM t0 JOIN t1 ON t0.a < t1.a", + "SELECT t0.a FROM t0 JOIN t1 ON t1.a > t0.a", + "SELECT t0.a FROM t0 JOIN t1 ON t0.a + 1 < t1.a * 2", + "SELECT t0.a FROM t0 WHERE EXISTS (SELECT 1 FROM t1 WHERE t0.a > t1.a)", + "SELECT t0.a FROM t0 WHERE NOT EXISTS (SELECT 1 FROM t1 WHERE t0.a > t1.a)", + ] { + let plan = ctx.sql(query).await?.create_physical_plan().await?; + let mut found = false; + plan.apply(|node| { + found |= node.downcast_ref::().is_some(); + Ok(TreeNodeRecursion::Continue) + })?; + assert!(found, "expected a PiecewiseMergeJoinExec for: {query}"); + + roundtrip_test(plan)?; + } + Ok(()) +}