From d488899caa0c61ae8f48c9939d971d6063b69f30 Mon Sep 17 00:00:00 2001 From: Geoffrey Claude Date: Fri, 28 Aug 2026 16:54:06 +0200 Subject: [PATCH 1/2] fix(examples): apply relation aliases once --- .../examples/relation_planner/table_sample.rs | 44 ++++++++++++++++++- datafusion/expr/src/planner.rs | 19 +++++++- .../library-user-guide/extending-sql.md | 7 +++ 3 files changed, 66 insertions(+), 4 deletions(-) diff --git a/datafusion-examples/examples/relation_planner/table_sample.rs b/datafusion-examples/examples/relation_planner/table_sample.rs index 7a8f533ac3a9b..b688402172aa4 100644 --- a/datafusion-examples/examples/relation_planner/table_sample.rs +++ b/datafusion-examples/examples/relation_planner/table_sample.rs @@ -373,10 +373,13 @@ impl RelationPlanner for TableSamplePlanner { }) .transpose()?; - // Plan the underlying table without the sample clause + // Plan the underlying table without the sample clause or alias. The + // alias belongs to the complete TABLESAMPLE relation, so we return it + // with `PlannedRelation` below and let DataFusion apply it once, after + // the sampling node has been added. let base_relation = TableFactor::Table { sample: None, - alias: alias.clone(), + alias: None, name, args, with_hints, @@ -467,6 +470,43 @@ impl RelationPlanner for TableSamplePlanner { } } +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn table_sample_applies_relation_alias_once() -> Result<()> { + let ctx = SessionContext::new(); + ctx.register_relation_planner(Arc::new(TableSamplePlanner))?; + register_sample_data(&ctx)?; + + let df = ctx + .sql( + "SELECT * FROM sample_data AS sampled(first, second) \ + TABLESAMPLE (3 ROWS)", + ) + .await?; + let plan = df.logical_plan().display_indent().to_string(); + + assert_snapshot!(plan, @r" + Projection: sampled.first, sampled.second + SubqueryAlias: sampled + Projection: sample_data.column1 AS first, sample_data.column2 AS second + Limit: skip=0, fetch=3 + TableScan: sample_data + "); + + // Keep the intent of this regression test obvious even if the plan + // display gains unrelated detail: the relation alias and its column + // rename projection each belong around the sampled relation once. + assert_eq!(plan.matches("SubqueryAlias: sampled").count(), 1); + assert_eq!(plan.matches(" AS first").count(), 1); + assert_eq!(plan.matches(" AS second").count(), 1); + + Ok(()) + } +} + /// Custom logical plan node representing a TABLESAMPLE operation. /// /// Stores sampling parameters (bounds, seed) and wraps the input plan. diff --git a/datafusion/expr/src/planner.rs b/datafusion/expr/src/planner.rs index 7aaf3a98cbe5d..be7e869e3c458 100644 --- a/datafusion/expr/src/planner.rs +++ b/datafusion/expr/src/planner.rs @@ -359,13 +359,22 @@ pub enum PlannerResult { pub struct PlannedRelation { /// The logical plan for the relation pub plan: LogicalPlan, - /// Optional table alias for the relation + /// Optional alias for the complete relation. + /// + /// DataFusion applies this alias after the extension planner's logical + /// plan has been created. If the extension recursively plans an inner + /// [`TableFactor`] with [`RelationPlannerContext::plan`], it should remove + /// this alias from the inner relation and return it here instead. Otherwise + /// the same alias may be applied to both the inner and outer plans. pub alias: Option, } #[cfg(feature = "sql")] impl PlannedRelation { - /// Create a new `PlannedRelation` with the given plan and alias + /// Create a new `PlannedRelation` with the given plan and optional alias. + /// + /// The plan should not already contain this alias. DataFusion applies it + /// after the relation planner returns. pub fn new(plan: LogicalPlan, alias: Option) -> Self { Self { plan, alias } } @@ -414,6 +423,12 @@ pub trait RelationPlannerContext { /// Plans the specified relation through the full planner pipeline, starting /// from the first registered relation planner. + /// + /// This also applies any alias present on `relation`. When unwrapping a + /// relation such as `TABLESAMPLE`, remove the outer relation's alias before + /// calling this method, then return that alias as part of the final + /// [`PlannedRelation`]. This keeps the alias around the complete relation + /// and ensures it is applied exactly once. fn plan(&mut self, relation: TableFactor) -> Result; /// Converts a SQL expression into a logical expression using the current diff --git a/docs/source/library-user-guide/extending-sql.md b/docs/source/library-user-guide/extending-sql.md index eea5b3b1acfc9..952aa5559b310 100644 --- a/docs/source/library-user-guide/extending-sql.md +++ b/docs/source/library-user-guide/extending-sql.md @@ -285,6 +285,12 @@ There are two main approaches when implementing a [`RelationPlanner`]: represent the operation in the logical plan, along with a custom [`ExecutionPlan`] to execute it. Both are required for end-to-end execution. +When returning a [`PlannedRelation`], pass along the alias from the relation and +let DataFusion apply it to the finished plan. If your planner removes a wrapper +or modifier and recursively calls `ctx.plan(...)`, remove that outer alias from +the inner relation first. Otherwise, the recursive call applies the alias once +and DataFusion applies it a second time when your planner returns. + #### Example: Basic RelationPlanner Structure ```rust @@ -380,6 +386,7 @@ SELECT * FROM sales [`sessioncontext`]: https://docs.rs/datafusion/latest/datafusion/execution/context/struct.SessionContext.html [`sessionstatebuilder`]: https://docs.rs/datafusion/latest/datafusion/execution/session_state/struct.SessionStateBuilder.html [`relationplannercontext`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/planner/trait.RelationPlannerContext.html +[`plannedrelation`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/planner/struct.PlannedRelation.html [exprplanner api documentation]: https://docs.rs/datafusion/latest/datafusion/logical_expr/planner/trait.ExprPlanner.html [typeplanner api documentation]: https://docs.rs/datafusion/latest/datafusion/logical_expr/planner/trait.TypePlanner.html [relationplanner api documentation]: https://docs.rs/datafusion/latest/datafusion/logical_expr/planner/trait.RelationPlanner.html From 9e3f228dba0f046b9a7c5ddfdfbdd384f28be933 Mon Sep 17 00:00:00 2001 From: Geoffrey Claude Date: Fri, 28 Aug 2026 16:54:41 +0200 Subject: [PATCH 2/2] docs: update boxed RelationPlanning example RelationPlanning::Planned and RelationPlanning::Original began taking boxed values in https://github.com/apache/datafusion/commit/e8efd59203283facc0a87b3afb9f5d2d909cc75e, merged through https://github.com/apache/datafusion/pull/19672. The extending SQL guide retained the old constructors, so update the example to match the current API. --- docs/source/library-user-guide/extending-sql.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/docs/source/library-user-guide/extending-sql.md b/docs/source/library-user-guide/extending-sql.md index 952aa5559b310..8899e4db2b364 100644 --- a/docs/source/library-user-guide/extending-sql.md +++ b/docs/source/library-user-guide/extending-sql.md @@ -320,11 +320,13 @@ impl RelationPlanner for MyRelationPlanner { // Transform or wrap the plan as needed // ... - Ok(RelationPlanning::Planned(PlannedRelation::new(input, alias))) + Ok(RelationPlanning::Planned(Box::new(PlannedRelation::new( + input, alias, + )))) } // Return Original for relations you don't handle - other => Ok(RelationPlanning::Original(other)), + other => Ok(RelationPlanning::Original(Box::new(other))), } } }