Skip to content

Commit e711398

Browse files
adriangbclaude
andcommitted
feat: protobuf serialization for the materialized CTE plan nodes
Add protobuf encodings for LogicalPlan::MaterializedCte, LogicalPlan::MaterializedCteScan, MaterializedCteExec and MaterializedCteScanExec. The physical encoding carries only the id of the CTE. The decode context holds one buffer for each id, so a decoded scan and its MaterializedCteExec share the buffer wherever they are in the plan, also when the scan is in a scalar subquery. Decoded buffers get new ids, because the ids of another process can be equal to local ids. A decoded plan fragment that has a scan without its MaterializedCteExec fails with a clear error when the scan is executed. Closes #26231. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 parent 1ca8c39 commit e711398

12 files changed

Lines changed: 1297 additions & 11 deletions

File tree

‎datafusion/expr/src/logical_plan/materialized_cte.rs‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,17 @@ impl MaterializedCteId {
5757
pub fn as_u64(self) -> u64 {
5858
self.0
5959
}
60+
61+
/// Create an id from the value that [`Self::as_u64`] returned, to decode
62+
/// a serialized plan.
63+
///
64+
/// Such an id comes from another process, so it can be equal to an id
65+
/// that [`Self::next`] returns in this process. This is safe in a logical
66+
/// plan: a scan reads the nearest [`MaterializedCte`] above it that has
67+
/// its id, and the physical planner gives every CTE a new id.
68+
pub fn from_u64(id: u64) -> Self {
69+
Self(id)
70+
}
6071
}
6172

6273
impl fmt::Display for MaterializedCteId {

‎datafusion/physical-plan/src/materialized_cte.rs‎

Lines changed: 203 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
//! disk is written in pieces, so that the move does not raise the memory that
4141
//! each reader needs.
4242
43+
use std::collections::HashMap;
4344
use std::fmt;
4445
use std::pin::Pin;
4546
use std::sync::Arc;
@@ -49,10 +50,11 @@ use arrow::datatypes::SchemaRef;
4950
use arrow::record_batch::RecordBatch;
5051
use datafusion_common::config::ConfigOptions;
5152
use datafusion_common::tree_node::TreeNodeRecursion;
52-
use datafusion_common::{Result, Statistics, internal_err};
53+
use datafusion_common::{Result, Statistics, exec_err, internal_err};
5354
use datafusion_common_runtime::JoinSet;
5455
use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation};
5556
use datafusion_execution::{SendableRecordBatchStream, SpillFile, TaskContext};
57+
use datafusion_expr::MaterializedCteId;
5658
use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr};
5759
use futures::stream::BoxStream;
5860
use futures::{Stream, StreamExt, TryStreamExt};
@@ -300,6 +302,11 @@ impl MaterializedCteBuffer {
300302
self.id
301303
}
302304

305+
/// The name of the CTE.
306+
pub fn name(&self) -> &str {
307+
&self.name
308+
}
309+
303310
fn set_body(&self, body: Arc<dyn ExecutionPlan>) {
304311
*self.body.lock() = Some(body);
305312
}
@@ -314,8 +321,13 @@ impl MaterializedCteBuffer {
314321
self: &Arc<Self>,
315322
context: Arc<TaskContext>,
316323
) -> Result<OnceFut<MaterializedOutput>> {
324+
// A plan from the planner always has the body. A decoded plan does
325+
// not have it when the plan was split between the body and a scan.
317326
let Some(body) = self.body.lock().clone() else {
318-
return internal_err!("MaterializedCte {} has no body", self.name);
327+
return exec_err!(
328+
"The materialized CTE {} has a MaterializedCteScanExec but no MaterializedCteExec in this plan. The body and the scans of a materialized CTE must be in one plan and run in one process",
329+
self.name
330+
);
319331
};
320332
let once = Arc::clone(&self.output.lock());
321333
let name = self.name.clone();
@@ -327,6 +339,36 @@ impl MaterializedCteBuffer {
327339
}
328340
}
329341

342+
/// The buffers of the materialized CTEs in one decoded plan.
343+
///
344+
/// A [`MaterializedCteExec`] and its [`MaterializedCteScanExec`]s are encoded
345+
/// as separate nodes that have the same id. The decoder of each node asks for
346+
/// the buffer of that id, so all of them share one buffer, in any decode
347+
/// order. A scan in a scalar subquery is not below its `MaterializedCteExec`,
348+
/// so one registry must be used for the whole plan.
349+
///
350+
/// The ids in an encoded plan come from another process and can be equal to
351+
/// the ids of plans made in this process. Each buffer therefore gets a new id
352+
/// from [`MaterializedCteId::next`].
353+
#[derive(Debug, Default)]
354+
pub struct DecodedMaterializedCteBuffers {
355+
/// The buffers, by the id in the encoded plan.
356+
buffers: Mutex<HashMap<u64, Arc<MaterializedCteBuffer>>>,
357+
}
358+
359+
impl DecodedMaterializedCteBuffers {
360+
/// Return the buffer of the CTE that has the id `encoded_id` in the
361+
/// encoded plan. The first call for an id creates the buffer.
362+
pub fn buffer(&self, encoded_id: u64, name: &str) -> Arc<MaterializedCteBuffer> {
363+
let mut buffers = self.buffers.lock();
364+
let buffer = buffers.entry(encoded_id).or_insert_with(|| {
365+
let id = MaterializedCteId::next().as_u64();
366+
Arc::new(MaterializedCteBuffer::new(id, name))
367+
});
368+
Arc::clone(buffer)
369+
}
370+
}
371+
330372
async fn buffer_body(
331373
name: String,
332374
body: Arc<dyn ExecutionPlan>,
@@ -655,6 +697,81 @@ impl ExecutionPlan for MaterializedCteExec {
655697
// continuation handles it.
656698
Ok(FilterPushdownPropagation::if_any(child_pushdown_result))
657699
}
700+
701+
#[cfg(feature = "proto")]
702+
fn try_to_proto(
703+
&self,
704+
ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>,
705+
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
706+
use datafusion_proto_models::protobuf;
707+
708+
// Destructure exhaustively (no `..`) so that adding a field to
709+
// `MaterializedCteExec` is a compile error here until it is either
710+
// serialized or explicitly documented as not needing to be.
711+
let Self {
712+
body,
713+
continuation,
714+
// Only the id and the name are encoded. The decoder makes a new
715+
// buffer and shares it with the scans that have the same id.
716+
buffer,
717+
// Copied from the continuation's properties, recomputed on decode.
718+
cache: _,
719+
} = self;
720+
Ok(Some(protobuf::PhysicalPlanNode {
721+
physical_plan_type: Some(
722+
protobuf::physical_plan_node::PhysicalPlanType::MaterializedCte(
723+
Box::new(protobuf::MaterializedCteExecNode {
724+
id: buffer.id(),
725+
name: buffer.name().to_string(),
726+
body: Some(Box::new(ctx.encode_child(body)?)),
727+
continuation: Some(Box::new(ctx.encode_child(continuation)?)),
728+
}),
729+
),
730+
),
731+
}))
732+
}
733+
}
734+
735+
#[cfg(feature = "proto")]
736+
impl MaterializedCteExec {
737+
/// Reconstruct a [`MaterializedCteExec`] from its protobuf representation.
738+
///
739+
/// The node shares the buffer that the decode context holds for its id
740+
/// with the decoded [`MaterializedCteScanExec`]s of the same id.
741+
pub fn try_from_proto(
742+
node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
743+
ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
744+
) -> Result<Arc<dyn ExecutionPlan>> {
745+
use datafusion_proto_models::protobuf;
746+
747+
let cte = crate::expect_plan_variant!(
748+
node,
749+
protobuf::physical_plan_node::PhysicalPlanType::MaterializedCte,
750+
"MaterializedCteExec",
751+
);
752+
// Destructure exhaustively so that a new field on
753+
// `MaterializedCteExecNode` is a compile error here rather than a
754+
// silently dropped field.
755+
let protobuf::MaterializedCteExecNode {
756+
id,
757+
name,
758+
body,
759+
continuation,
760+
} = &**cte;
761+
let body =
762+
ctx.decode_required_child(body.as_deref(), "MaterializedCteExec", "body")?;
763+
let continuation = ctx.decode_required_child(
764+
continuation.as_deref(),
765+
"MaterializedCteExec",
766+
"continuation",
767+
)?;
768+
let buffer = ctx.materialized_cte_buffer(*id, name)?;
769+
// A second node would replace the body that the scans read.
770+
if buffer.body.lock().is_some() {
771+
return internal_err!("Two MaterializedCteExec nodes have the id {id}");
772+
}
773+
Ok(Arc::new(Self::new(body, continuation, buffer)))
774+
}
658775
}
659776

660777
/// Reads the buffered output of a [`MaterializedCteExec`] body.
@@ -815,6 +932,90 @@ impl ExecutionPlan for MaterializedCteScanExec {
815932
) -> Result<Arc<Statistics>> {
816933
Ok(Arc::new(Statistics::new_unknown(&self.schema)))
817934
}
935+
936+
#[cfg(feature = "proto")]
937+
fn try_to_proto(
938+
&self,
939+
_ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>,
940+
) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
941+
use datafusion_common::utils::usize_to_wire;
942+
use datafusion_proto_models::protobuf;
943+
944+
// Destructure exhaustively (no `..`) so that adding a field to
945+
// `MaterializedCteScanExec` is a compile error here until it is
946+
// either serialized or explicitly documented as not needing to be.
947+
let Self {
948+
id,
949+
name,
950+
schema,
951+
// The decoder binds the scan to the buffer of `id` again.
952+
buffer: _,
953+
// Runtime metrics, created empty on decode.
954+
metrics: _,
955+
cache,
956+
} = self;
957+
let partitions = usize_to_wire(
958+
cache.partitioning.partition_count(),
959+
"MaterializedCteScanExec",
960+
"partitions",
961+
)?;
962+
Ok(Some(protobuf::PhysicalPlanNode {
963+
physical_plan_type: Some(
964+
protobuf::physical_plan_node::PhysicalPlanType::MaterializedCteScan(
965+
protobuf::MaterializedCteScanExecNode {
966+
id: *id,
967+
name: name.clone(),
968+
schema: Some(schema.as_ref().try_into()?),
969+
partitions,
970+
},
971+
),
972+
),
973+
}))
974+
}
975+
}
976+
977+
#[cfg(feature = "proto")]
978+
impl MaterializedCteScanExec {
979+
/// Reconstruct a [`MaterializedCteScanExec`] from its protobuf
980+
/// representation.
981+
///
982+
/// The scan is bound to the buffer that the decode context holds for its
983+
/// id. The decoded [`MaterializedCteExec`] of the same id gives that
984+
/// buffer its body. If the decoded plan has no such node, the scan fails
985+
/// when it is executed.
986+
pub fn try_from_proto(
987+
node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
988+
ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
989+
) -> Result<Arc<dyn ExecutionPlan>> {
990+
use datafusion_common::utils::usize_from_wire;
991+
use datafusion_proto_models::protobuf;
992+
993+
let scan = crate::expect_plan_variant!(
994+
node,
995+
protobuf::physical_plan_node::PhysicalPlanType::MaterializedCteScan,
996+
"MaterializedCteScanExec",
997+
);
998+
// Destructure exhaustively so that a new field on
999+
// `MaterializedCteScanExecNode` is a compile error here rather than a
1000+
// silently dropped field.
1001+
let protobuf::MaterializedCteScanExecNode {
1002+
id,
1003+
name,
1004+
schema,
1005+
partitions,
1006+
} = scan;
1007+
let schema = schema.as_ref().ok_or_else(|| {
1008+
datafusion_common::internal_datafusion_err!(
1009+
"MaterializedCteScanExec is missing required field 'schema'"
1010+
)
1011+
})?;
1012+
let schema = Arc::new(arrow::datatypes::Schema::try_from(schema)?);
1013+
let partitions =
1014+
usize_from_wire(*partitions, "MaterializedCteScanExec", "partitions")?;
1015+
let buffer = ctx.materialized_cte_buffer(*id, name)?;
1016+
let scan = Self::new(buffer.id(), name, schema, partitions);
1017+
Ok(Arc::new(scan.bind(buffer)))
1018+
}
8181019
}
8191020

8201021
#[cfg(test)]

‎datafusion/physical-plan/src/proto/mod.rs‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,7 @@ use datafusion_proto_models::protobuf::{
9191
};
9292

9393
use crate::ExecutionPlan;
94+
use crate::materialized_cte::MaterializedCteBuffer;
9495

9596
pub use datafusion_proto_models::ProtoDecoderRegistry;
9697
// Not public API: see `decode_execution_plan`.
@@ -181,6 +182,16 @@ pub trait ExecutionPlanDecode {
181182

182183
/// Reconstruct a window UDF from its name and optional payload.
183184
fn decode_udwf(&self, name: &str, payload: Option<&[u8]>) -> Result<Arc<WindowUDF>>;
185+
186+
/// The buffer of the materialized CTE that has the id `encoded_id` in the
187+
/// encoded plan. Every call for one id during the decoding of one plan
188+
/// must return the same buffer. See
189+
/// [`DecodedMaterializedCteBuffers`](crate::materialized_cte::DecodedMaterializedCteBuffers).
190+
fn materialized_cte_buffer(
191+
&self,
192+
encoded_id: u64,
193+
name: &str,
194+
) -> Result<Arc<MaterializedCteBuffer>>;
184195
}
185196

186197
/// Context handed to [`ExecutionPlan::try_to_proto`].
@@ -414,6 +425,16 @@ impl<'a> ExecutionPlanDecodeCtx<'a> {
414425
self.decoder.decode_udwf(name, payload)
415426
}
416427

428+
/// The buffer that a decoded `MaterializedCteExec` shares with the decoded
429+
/// `MaterializedCteScanExec`s that have the same `encoded_id`.
430+
pub fn materialized_cte_buffer(
431+
&self,
432+
encoded_id: u64,
433+
name: &str,
434+
) -> Result<Arc<MaterializedCteBuffer>> {
435+
self.decoder.materialized_cte_buffer(encoded_id, name)
436+
}
437+
417438
/// An expression-level decode context backed by this plan context, bound to
418439
/// `input_schema`.
419440
///

‎datafusion/physical-plan/src/proto_test_util.rs‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ use datafusion_proto_models::protobuf::{
5454

5555
use crate::ExecutionPlan;
5656
use crate::empty::EmptyExec;
57+
use crate::materialized_cte::MaterializedCteBuffer;
5758
use crate::proto::{ExecutionPlanDecode, ExecutionPlanEncode};
5859

5960
/// The schema shared by the stub child plans: `a: Int32, b: Int32`.
@@ -318,6 +319,14 @@ impl ExecutionPlanDecode for StubPlanDecoder {
318319
fn decode_udwf(&self, name: &str, _payload: Option<&[u8]>) -> Result<Arc<WindowUDF>> {
319320
internal_err!("stub decoder cannot decode the window UDF {name}")
320321
}
322+
323+
fn materialized_cte_buffer(
324+
&self,
325+
_encoded_id: u64,
326+
name: &str,
327+
) -> Result<Arc<MaterializedCteBuffer>> {
328+
internal_err!("stub decoder cannot decode the materialized CTE {name}")
329+
}
321330
}
322331

323332
/// Minimal extension plans for the [`crate::proto::registry`] tests.
@@ -526,4 +535,14 @@ impl ExecutionPlanDecode for UnreachablePlanDecoder {
526535
) -> Result<Arc<WindowUDF>> {
527536
unreachable!("decode_udwf must not be reached when the node is rejected")
528537
}
538+
539+
fn materialized_cte_buffer(
540+
&self,
541+
_encoded_id: u64,
542+
_name: &str,
543+
) -> Result<Arc<MaterializedCteBuffer>> {
544+
unreachable!(
545+
"materialized_cte_buffer must not be reached when the node is rejected"
546+
)
547+
}
529548
}

0 commit comments

Comments
 (0)