asolimando commented on code in PR #26428:
URL: https://github.com/apache/datafusion/pull/26428#discussion_r4437457791
##########
datafusion/physical-plan/src/materialized_cte.rs:
##########
@@ -816,10 +929,20 @@ impl ExecutionPlan for MaterializedCteScanExec {
);
};
let mut output = buffer.materialize(context)?;
- let output_partitions = self.cache.partitioning.partition_count();
+ let output_partitions = self.partition_count();
+ let mirrors_body = self.mirrors_body;
+ let name = self.name.clone();
let baseline = BaselineMetrics::new(&self.metrics, partition);
let stream = futures::stream::once(async move {
let output = std::future::poll_fn(|cx|
output.get_shared(cx)).await?;
+ // The ordering and the hash partitioning of the scan are only
+ // correct when one scan partition replays one body partition.
+ if mirrors_body && output.partitions.len() != output_partitions {
Review Comment:
Here we check only the partition count, but if you have a body with the same
count but a different ordering or hash keys, it goes through with wrong
results. I have checked with a variant of
`mirror_scan_rejects_other_body_partition_count`: same buffer, unsorted body,
the scan still reports `[a ASC]` and returns `[3, 1, 2]`.
I don't think the problem can reproduce in DF core, but projects doing
post-planning rewrites like Ballista or datafusion distributed could, we should
probably handle this to be on the safe side.
A cheap fix would be checking the buffer's body in `execute` has the same
hash partitioning, and an ordering that satisfies the scan's.
If you prefer not to fix it, I think we should at least put a note on
`MaterializedCteExec::new`'s documentation saying the body must keep the copied
properties.
WDYT?
##########
datafusion/physical-plan/src/materialized_cte.rs:
##########
Review Comment:
This points to the issue this PR closes, but statistics are still only known
for single-partition scans. Shall we file a new issue and repoint it?
##########
datafusion/core/src/physical_planner.rs:
##########
@@ -317,125 +321,344 @@ fn renumber_materialized_ctes(plan: &LogicalPlan) ->
Result<Option<LogicalPlan>>
}
/// Points every [`MaterializedCteScanExec`] at the buffer of the
-/// [`MaterializedCteExec`] with the same id, and gives it the statistics of
-/// the CTE body.
+/// [`MaterializedCteExec`] with the same id.
///
/// [`renumber_materialized_ctes`] makes the ids unique before the physical
/// plan is created.
///
-/// A scan that is already bound was planned and bound by a separate planning
-/// run, such as the `scan` of a [`ViewTable`], and is not bound again. Binding
-/// it again would count its partitions twice in the buffer.
+/// A scan that is already bound comes from a separate planning run, such as
+/// the `scan` of a [`ViewTable`], and is bound again. That run can have given
+/// it the ordering and the partitioning of a body that is optimized again
+/// now, and the new scan does not have them.
+/// [`register_materialized_ctes`] sets the partition counts of the scans of
+/// the final plan in the buffer.
///
/// Scans are planned before the node that owns their CTE, and a scan inside a
/// scalar subquery is planned in a separate subtree, so the binding is done
/// on the whole initial plan.
///
/// [`ViewTable`]: crate::datasource::ViewTable
-struct MaterializedCteBinder {
- ctes: HashMap<u64, Arc<dyn ExecutionPlan>>,
- /// The statistics of the CTE bodies. `None` while they are computed.
- statistics: HashMap<u64, Option<Arc<Statistics>>>,
- context: StatisticsContext,
+fn bind_materialized_cte_scans(
+ plan: Arc<dyn ExecutionPlan>,
+ ctes: &HashMap<u64, Arc<dyn ExecutionPlan>>,
+) -> Result<Arc<dyn ExecutionPlan>> {
+ plan.transform_up(|node| {
+ let Some(scan) = node.downcast_ref::<MaterializedCteScanExec>() else {
+ return Ok(Transformed::no(node));
+ };
+ let Some(cte) = ctes.get(&scan.id()) else {
+ if scan.is_bound() {
+ return Ok(Transformed::no(node));
+ }
+ return internal_err!("MaterializedCteScanExec {} has no CTE",
scan.id());
+ };
+ let scan = scan.bind(Arc::clone(as_materialized_cte(cte)?.buffer()));
+ Ok(Transformed::yes(Arc::new(scan) as Arc<dyn ExecutionPlan>))
+ })
+ .map(|t| t.data)
}
-impl MaterializedCteBinder {
- fn new(plan: &Arc<dyn ExecutionPlan>, session_state: &dyn Session) ->
Result<Self> {
- let mut ctes = HashMap::new();
- plan.apply(|node| {
- if let Some(cte) = node.downcast_ref::<MaterializedCteExec>()
- && ctes.insert(cte.buffer().id(), Arc::clone(node)).is_some()
+/// The [`MaterializedCteExec`] nodes of `plan`, by id.
+fn materialized_ctes(
+ plan: &Arc<dyn ExecutionPlan>,
+) -> Result<HashMap<u64, Arc<dyn ExecutionPlan>>> {
+ let mut ctes = HashMap::new();
+ plan.apply(|node| {
+ if let Some(cte) = node.downcast_ref::<MaterializedCteExec>()
+ && ctes.insert(cte.buffer().id(), Arc::clone(node)).is_some()
+ {
+ return internal_err!(
+ "Two MaterializedCteExec nodes have the id {}",
+ cte.buffer().id()
+ );
+ }
+ Ok(TreeNodeRecursion::Continue)
+ })?;
+ Ok(ctes)
+}
+
+fn as_materialized_cte(plan: &Arc<dyn ExecutionPlan>) ->
Result<&MaterializedCteExec> {
+ match plan.downcast_ref::<MaterializedCteExec>() {
+ Some(cte) => Ok(cte),
+ None => internal_err!("Expected a MaterializedCteExec"),
+ }
+}
+
+/// The ids of `ctes`, in an order in which each CTE comes after the CTEs
+/// that its body reads or defines.
+fn materialized_cte_order(
+ ctes: &HashMap<u64, Arc<dyn ExecutionPlan>>,
+) -> Result<Vec<u64>> {
+ fn visit(
+ id: u64,
+ ctes: &HashMap<u64, Arc<dyn ExecutionPlan>>,
+ visiting: &mut HashSet<u64>,
+ order: &mut Vec<u64>,
+ ) -> Result<()> {
+ if order.contains(&id) {
+ return Ok(());
+ }
+ if !visiting.insert(id) {
+ return internal_err!("The body of MaterializedCteExec {id} reads
itself");
+ }
+ let mut reads = vec![];
+ as_materialized_cte(&ctes[&id])?.body().apply(|node| {
+ if let Some(scan) = node.downcast_ref::<MaterializedCteScanExec>()
+ && ctes.contains_key(&scan.id())
{
- return internal_err!(
- "Two MaterializedCteExec nodes have the id {}",
- cte.buffer().id()
- );
+ reads.push(scan.id());
+ }
+ if let Some(cte) = node.downcast_ref::<MaterializedCteExec>() {
+ reads.push(cte.buffer().id());
}
Ok(TreeNodeRecursion::Continue)
})?;
- let context = match session_state.statistics_registry() {
- Some(registry) =>
StatisticsContext::new_with_registry(registry.clone()),
- None => StatisticsContext::new(),
- };
- Ok(Self {
- ctes,
- statistics: HashMap::new(),
- context,
- })
+ for read in reads {
+ visit(read, ctes, visiting, order)?;
+ }
+ order.push(id);
+ Ok(())
}
- fn cte(&self, id: u64) -> Option<&MaterializedCteExec> {
- self.ctes
- .get(&id)
- .and_then(|cte| cte.downcast_ref::<MaterializedCteExec>())
+ let mut ids: Vec<u64> = ctes.keys().copied().collect();
+ ids.sort_unstable();
+ let mut visiting = HashSet::new();
+ let mut order = vec![];
+ for id in ids {
+ visit(id, ctes, &mut visiting, &mut order)?;
}
+ Ok(order)
+}
- /// Give every scan in `plan` whose CTE is in `self.ctes` the statistics
- /// of the CTE body. If `bind` is true, also bind the scans that are not
- /// bound.
- fn bind(
- &mut self,
- plan: Arc<dyn ExecutionPlan>,
- bind: bool,
- ) -> Result<Arc<dyn ExecutionPlan>> {
- plan.transform_up(|node| {
- let Some(scan) = node.downcast_ref::<MaterializedCteScanExec>()
else {
- return Ok(Transformed::no(node));
- };
- let Some(cte) = self.cte(scan.id()) else {
- if scan.is_bound() {
- return Ok(Transformed::no(node));
+/// A CTE body after the optimizers ran on it.
+struct OptimizedCteBody {
+ plan: Arc<dyn ExecutionPlan>,
+ statistics: Arc<Statistics>,
+}
+
+/// Prepare `plan` for the optimizers, after the bodies in `bodies` are
+/// optimized:
+///
+/// - Each scan whose CTE is in `bodies` gets the statistics, the partition
+/// count, the output ordering and the hash partitioning of that body.
+/// - Each [`MaterializedCteExec`] whose body is in `bodies` becomes a
+/// [`FrozenMaterializedCteExec`], so that the optimizers do not see or
+/// change the body again.
+fn freeze_materialized_ctes(
+ plan: Arc<dyn ExecutionPlan>,
+ bodies: &HashMap<u64, OptimizedCteBody>,
+) -> Result<Arc<dyn ExecutionPlan>> {
+ plan.transform_up(|node| {
+ if let Some(scan) = node.downcast_ref::<MaterializedCteScanExec>()
+ && let Some(body) = bodies.get(&scan.id())
+ {
+ let scan = scan.with_body_statistics(Arc::clone(&body.statistics));
+ let scan = scan
+ .with_body_properties(body.plan.as_ref())
+ .unwrap_or(scan);
+ return Ok(Transformed::yes(Arc::new(scan) as Arc<dyn
ExecutionPlan>));
+ }
+ if let Some(cte) = node.downcast_ref::<MaterializedCteExec>()
+ && let Some(body) = bodies.get(&cte.buffer().id())
+ {
+ let continuation = Arc::clone(node.children()[1]);
+ let frozen = FrozenMaterializedCteExec::new(
+ Arc::clone(&body.plan),
+ continuation,
+ Arc::clone(cte.buffer()),
+ );
+ return Ok(Transformed::yes(Arc::new(frozen) as Arc<dyn
ExecutionPlan>));
+ }
+ Ok(Transformed::no(node))
+ })
+ .map(|t| t.data)
+}
+
+/// Turn each [`FrozenMaterializedCteExec`] of `plan` back into a
+/// [`MaterializedCteExec`].
+fn thaw_materialized_ctes(
+ plan: Arc<dyn ExecutionPlan>,
+) -> Result<Arc<dyn ExecutionPlan>> {
+ plan.transform_up(|node| {
+ let Some(frozen) = node.downcast_ref::<FrozenMaterializedCteExec>()
else {
+ return Ok(Transformed::no(node));
+ };
+ let cte = MaterializedCteExec::new(
+ Arc::clone(&frozen.body),
+ Arc::clone(&frozen.continuation),
+ Arc::clone(&frozen.buffer),
+ );
+ Ok(Transformed::yes(Arc::new(cte) as Arc<dyn ExecutionPlan>))
+ })
+ .map(|t| t.data)
+}
+
+mod frozen_cte {
+ use std::fmt;
+ use std::sync::Arc;
+
+ use datafusion_common::config::ConfigOptions;
+ use datafusion_common::tree_node::TreeNodeRecursion;
+ use datafusion_common::{Result, Statistics, internal_err};
+ use datafusion_execution::{SendableRecordBatchStream, TaskContext};
+ use datafusion_physical_expr::PhysicalExpr;
+ use datafusion_physical_plan::execution_plan::{CardinalityEffect,
PlanProperties};
+ use datafusion_physical_plan::filter_pushdown::{
+ ChildFilterDescription, ChildPushdownResult, FilterDescription,
+ FilterPushdownPhase, FilterPushdownPropagation,
+ };
+ use datafusion_physical_plan::materialized_cte::MaterializedCteBuffer;
+ use datafusion_physical_plan::statistics::StatisticsArgs;
+ use datafusion_physical_plan::{DisplayAs, DisplayFormatType,
ExecutionPlan};
+
+ #[cfg(doc)]
+ use datafusion_physical_plan::materialized_cte::MaterializedCteExec;
+
+ /// A [`MaterializedCteExec`] whose body is already optimized, while the
+ /// optimizers run on the rest of the plan.
+ ///
+ /// The body is not a child, so the optimizers cannot change it, and the
+ /// properties that the scans copied from it stay true. For the
continuation,
+ /// this node acts like a [`MaterializedCteExec`].
+ #[derive(Debug)]
+ pub(super) struct FrozenMaterializedCteExec {
+ pub(super) body: Arc<dyn ExecutionPlan>,
+ pub(super) continuation: Arc<dyn ExecutionPlan>,
+ pub(super) buffer: Arc<MaterializedCteBuffer>,
+ }
+
+ impl FrozenMaterializedCteExec {
+ pub(super) fn new(
+ body: Arc<dyn ExecutionPlan>,
+ continuation: Arc<dyn ExecutionPlan>,
+ buffer: Arc<MaterializedCteBuffer>,
+ ) -> Self {
+ Self {
+ body,
+ continuation,
+ buffer,
+ }
+ }
+ }
+
+ impl DisplayAs for FrozenMaterializedCteExec {
+ fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) ->
fmt::Result {
+ match t {
+ DisplayFormatType::Default | DisplayFormatType::Verbose => {
+ write!(f, "MaterializedCteExec: name={}",
self.buffer.name())
Review Comment:
I understand we don't want to leak the freezing of CTEs in plans, but plan
display is the quickest debugging tool: in `EXPLAIN VERBOSE` the body
disappears after the first rule while the node still shows as a regular
`MaterializedCteExec`, so it looks like the body was dropped, which can be
confusing.
Is there a way we can mark the node as frozen, or still show its body,
somehow?
##########
datafusion/physical-plan/src/materialized_cte.rs:
##########
@@ -737,13 +786,67 @@ impl MaterializedCteScanExec {
id: self.id,
name: self.name.clone(),
schema: Arc::clone(&self.schema),
+ partitions: self.partitions,
+ mirrors_body: self.mirrors_body,
statistics: Some(statistics),
buffer: self.buffer.clone(),
metrics: ExecutionPlanMetricsSet::new(),
cache: Arc::clone(&self.cache),
}
}
+ /// Give the scan the partition count of `body`, its output ordering and
+ /// its hash partitioning. Partition `p` of the scan then replays
+ /// partition `p` of the body.
+ ///
+ /// Returns `None` when `body` has no output ordering and no hash
+ /// partitioning, or when its schema is not the schema of the scan.
+ ///
+ /// The caller must make sure that the body of the plan that is executed
+ /// is `body`, or has the same properties.
+ pub fn with_body_properties(&self, body: &dyn ExecutionPlan) ->
Option<Self> {
+ let body_schema = body.schema();
+ let same_columns = self.schema.fields().len() ==
body_schema.fields().len()
+ && self
+ .schema
+ .fields()
+ .iter()
+ .zip(body_schema.fields())
+ .all(|(a, b)| a.name() == b.name() && a.data_type() ==
b.data_type());
+ if !same_columns {
+ return None;
+ }
+ let ordering = body.output_ordering().cloned();
Review Comment:
Nit: this keeps only the first ordering, and drops the equivalences and
constants of the body, but they would still be valid for the scan. Copying the
body’s equivalence properties (with the scan schema) would preserve them.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
