You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Commit e711398
Browse filesBrowse the repository at this point in the historyBrowse files
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>
Copy file name to clipboardExpand all lines: datafusion/physical-plan/src/materialized_cte.rs
+203-2Lines changed: 203 additions & 2 deletions
Original file line number
Diff line number
Diff line change
@@ -40,6 +40,7 @@
40
40
//! disk is written in pieces, so that the move does not raise the memory that
41
41
//! each reader needs.
42
42
43
+
use std::collections::HashMap;
43
44
use std::fmt;
44
45
use std::pin::Pin;
45
46
use std::sync::Arc;
@@ -49,10 +50,11 @@ use arrow::datatypes::SchemaRef;
49
50
use arrow::record_batch::RecordBatch;
50
51
use datafusion_common::config::ConfigOptions;
51
52
use datafusion_common::tree_node::TreeNodeRecursion;
52
-
use datafusion_common::{Result,Statistics, internal_err};
53
+
use datafusion_common::{Result,Statistics,exec_err,internal_err};
53
54
use datafusion_common_runtime::JoinSet;
54
55
use datafusion_execution::memory_pool::{MemoryConsumer,MemoryReservation};
55
56
use datafusion_execution::{SendableRecordBatchStream,SpillFile,TaskContext};
57
+
use datafusion_expr::MaterializedCteId;
56
58
use datafusion_physical_expr::{EquivalenceProperties,Partitioning,PhysicalExpr};
57
59
use futures::stream::BoxStream;
58
60
use futures::{Stream,StreamExt,TryStreamExt};
@@ -300,6 +302,11 @@ impl MaterializedCteBuffer {
300
302
self.id
301
303
}
302
304
305
+
/// The name of the CTE.
306
+
pubfnname(&self) -> &str{
307
+
&self.name
308
+
}
309
+
303
310
fnset_body(&self,body:Arc<dynExecutionPlan>){
304
311
*self.body.lock() = Some(body);
305
312
}
@@ -314,8 +321,13 @@ impl MaterializedCteBuffer {
314
321
self:&Arc<Self>,
315
322
context:Arc<TaskContext>,
316
323
) -> 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.
317
326
letSome(body) = self.body.lock().clone()else{
318
-
returninternal_err!("MaterializedCte {} has no body",self.name);
327
+
returnexec_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
+
);
319
331
};
320
332
let once = Arc::clone(&self.output.lock());
321
333
let name = self.name.clone();
@@ -327,6 +339,36 @@ impl MaterializedCteBuffer {
327
339
}
328
340
}
329
341
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
0 commit comments