Skip to content

Commit e071ec8

Browse files
committed
feat: core SPI for contrib leaf scans (CometScanWithPlanData)
Introduce a small extension contract so out-of-tree Comet contrib leaf scans (Delta, and future Hudi/etc.) can participate in native planning without core holding a compile-time reference to them -- mirroring the Iceberg-precedent of keeping the data-source-specific code at the edge. What this adds: - `trait CometScanWithPlanData` (`sourceKey` / `commonData` / `perPartitionData`, plus optional `dynamicPruningFilters` / `withDynamicPruningFilters` for scans whose DPP filters live in a @transient field). `CometNativeScanExec` now mixes it in. - `CometNativeExec.foreachUntilCometInput` matches `case _: CometLeafExec` (a strict superset of the previous fixed scan enumeration -- all built-in leaf scans already extend `CometLeafExec`), so any leaf Comet exec is recognised as an input boundary. - `PlanDataInjector.findAllPlanData` collects per-partition planning data via the trait instead of a hardcoded `CometNativeScanExec` match. - `PlanDataInjector`'s registry gains one reflective `DeltaPlanDataInjector$` slot, appended only when a contrib bundled it (`-Pcontrib-delta`). Default builds get a `ClassNotFoundException` -> `None` and an unchanged injectors list, so there is zero contrib surface at runtime. - `CometPlanAdaptiveDynamicPruningFilters` rewrites AQE DPP filters in place for trait scans whose filters can't survive `makeCopy` (apache#3510). Inert by construction: with no contrib on the classpath this is behavior- preserving (the leaf match is a superset; the trait match catches the same `CometNativeScanExec`; the reflective slot resolves to nothing). Tests: `CometScanWithPlanDataSuite` (trait-contract defaults + reflective-slot graceful absence). Verified `CometJoinSuite` (native scan fusion / DPP) stays green. First unit of the Delta-contrib PR split (tracking: apache#4366).
1 parent 0031c60 commit e071ec8

4 files changed

Lines changed: 190 additions & 16 deletions

File tree

spark/src/main/scala/org/apache/comet/rules/CometPlanAdaptiveDynamicPruningFilters.scala

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,15 @@ case object CometPlanAdaptiveDynamicPruningFilters
8383
if icebergScan.runtimeFilters.exists(hasCometSAB) =>
8484
logDebug("Converting AQE DPP for CometIcebergNativeScanExec")
8585
convertIcebergScanDPP(icebergScan, plan)
86+
// Comet scans whose DPP filters live in a @transient field (the contrib's
87+
// CometDeltaNativeScanExec). transformExpressions/makeCopy can't rewrite
88+
// them, and a rewritten copy is orphaned when the enclosing native block
89+
// is rebuilt (#3510). The scan's `withDynamicPruningFilters` installs the
90+
// rewrite in place and returns `this`, so it lands on the executing
91+
// instance.
92+
case p: CometScanWithPlanData if p.dynamicPruningFilters.exists(hasCometSAB) =>
93+
logDebug(s"Converting AQE DPP for ${p.getClass.getSimpleName} in place")
94+
p.withDynamicPruningFilters(p.dynamicPruningFilters.map(f => convertFilter(f, plan)))
8695
case p: SparkPlan
8796
if !p.isInstanceOf[CometNativeScanExec]
8897
&& !p.isInstanceOf[CometIcebergNativeScanExec]

spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,8 @@ case class CometNativeScanExec(
7171
sourceKey: String) // Key for PlanDataInjector to match common+partition data at runtime
7272
extends CometLeafExec
7373
with DataSourceScanExec
74-
with ShimStreamSourceAwareSparkPlan {
74+
with ShimStreamSourceAwareSparkPlan
75+
with CometScanWithPlanData {
7576

7677
override lazy val metadata: Map[String, String] =
7778
if (originalPlan != null) originalPlan.metadata else Map.empty

spark/src/main/scala/org/apache/spark/sql/comet/operators.scala

Lines changed: 93 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import scala.jdk.CollectionConverters._
2727

2828
import org.apache.spark.{Partition, TaskContext}
2929
import org.apache.spark.broadcast.Broadcast
30+
import org.apache.spark.internal.Logging
3031
import org.apache.spark.rdd.RDD
3132
import org.apache.spark.sql.catalyst.InternalRow
3233
import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeSet, Expression, ExpressionSet, Generator, NamedExpression, SortOrder}
@@ -81,14 +82,41 @@ private[comet] trait PlanDataInjector {
8182
/**
8283
* Registry and utilities for injecting per-partition planning data into operator trees.
8384
*/
84-
private[comet] object PlanDataInjector {
85-
86-
// Registry of injectors for different operator types
87-
private val injectors: Seq[PlanDataInjector] = Seq(
88-
IcebergPlanDataInjector,
89-
NativeScanPlanDataInjector
90-
// Future: DeltaPlanDataInjector, HudiPlanDataInjector, etc.
91-
)
85+
private[comet] object PlanDataInjector extends Logging {
86+
87+
// Registry of injectors for different operator types. The contrib/delta integration's
88+
// DeltaPlanDataInjector is appended via one reflective class lookup -- present only when
89+
// the contrib was bundled (i.e. -Pcontrib-delta on the Maven build). Default builds get
90+
// the empty Option and an unmodified injectors list, so there's zero contrib surface at
91+
// runtime on default builds.
92+
private val injectors: Seq[PlanDataInjector] = {
93+
val builtin: Seq[PlanDataInjector] = Seq(IcebergPlanDataInjector, NativeScanPlanDataInjector)
94+
val deltaOpt: Option[PlanDataInjector] =
95+
try {
96+
// Scala compiles `object Foo` into BOTH `Foo.class` (a static-forwarder
97+
// class with no MODULE$ field) AND `Foo$.class` (the module class that
98+
// does have MODULE$). The trailing `$` selects the module class.
99+
// scalastyle:off classforname
100+
val cls = Class.forName("org.apache.spark.sql.comet.DeltaPlanDataInjector$")
101+
// scalastyle:on classforname
102+
Some(cls.getField("MODULE$").get(null).asInstanceOf[PlanDataInjector])
103+
} catch {
104+
// Default builds (no -Pcontrib-delta) won't have the class -> silent None.
105+
case _: ClassNotFoundException => None
106+
// Reflection-binding failures (signature/access drift) -> silent None.
107+
case _: NoSuchFieldException | _: IllegalAccessException => None
108+
// Anything else (ExceptionInInitializerError, linkage errors, CCE on the
109+
// PlanDataInjector cast) is a real bug -- surface it as a log warning and
110+
// still decline so the rest of the planner stays alive.
111+
case e: Throwable =>
112+
logWarning(
113+
"Found org.apache.spark.sql.comet.DeltaPlanDataInjector$ on classpath " +
114+
"but failed to load it; skipping contrib-delta plan-data injection",
115+
e)
116+
None
117+
}
118+
builtin ++ deltaOpt
119+
}
92120

93121
/**
94122
* Injects planning data into an Operator tree by finding nodes that need injection and applying
@@ -698,8 +726,16 @@ abstract class CometNativeExec extends CometExec {
698726
*/
699727
def foreachUntilCometInput(plan: SparkPlan)(func: SparkPlan => Unit): Unit = {
700728
plan match {
701-
case _: CometNativeScanExec | _: CometScanExec | _: CometBatchScanExec |
702-
_: CometIcebergNativeScanExec | _: CometCsvNativeScanExec | _: ShuffleQueryStageExec |
729+
// Match `CometLeafExec` first so contrib leaf scans (e.g. the Delta
730+
// contrib's `CometDeltaNativeScanExec`) are recognised as input boundaries
731+
// without requiring a core compile-time reference to the contrib class.
732+
// All built-in leaf scans (`CometNativeScanExec`, `CometIcebergNativeScanExec`,
733+
// `CometCsvNativeScanExec`) also extend `CometLeafExec`, so this is a
734+
// strict superset of the previous enumeration -- it just generalises the
735+
// input-boundary concept from "this fixed list" to "any leaf Comet exec".
736+
case _: CometLeafExec =>
737+
func(plan)
738+
case _: CometScanExec | _: CometBatchScanExec | _: ShuffleQueryStageExec |
703739
_: AQEShuffleReadExec | _: CometShuffleExchangeExec | _: CometUnionExec |
704740
_: CometTakeOrderedAndProjectExec | _: CometCoalesceExec | _: ReusedExchangeExec |
705741
_: CometBroadcastExchangeExec | _: BroadcastQueryStageExec |
@@ -766,11 +802,16 @@ abstract class CometNativeExec extends CometExec {
766802
(Map.empty, Map.empty)
767803
}
768804

769-
case nativeScan: CometNativeScanExec =>
770-
nativeScan.ensureSubqueriesResolved()
771-
(
772-
Map(nativeScan.sourceKey -> nativeScan.commonData),
773-
Map(nativeScan.sourceKey -> nativeScan.perPartitionData))
805+
// Generic path for leaf scans that surface planning data via the
806+
// `CometScanWithPlanData` trait. Catches `CometNativeScanExec` and any contrib
807+
// leaf scan (e.g. the Delta contrib's `CometDeltaNativeScanExec`) without
808+
// requiring core to compile-time reference contrib classes.
809+
case s: CometScanWithPlanData =>
810+
s match {
811+
case leaf: CometLeafExec => leaf.ensureSubqueriesResolved()
812+
case _ => // no DPP lifecycle to drive
813+
}
814+
(Map(s.sourceKey -> s.commonData), Map(s.sourceKey -> s.perPartitionData))
774815

775816
// Broadcast stages are boundaries - don't collect per-partition data from inside them.
776817
// After DPP filtering, broadcast scans may have different partition counts than the
@@ -881,6 +922,43 @@ abstract class CometLeafExec extends CometNativeExec with LeafExecNode {
881922
}
882923
}
883924

925+
/**
926+
* Marker trait for scan execs that surface planning data (a `commonData` block + per-partition
927+
* task bytes keyed by `sourceKey`) so that a parent `CometNativeExec` can find and inject the
928+
* data when the scan is fused into a larger native subtree.
929+
*
930+
* Implemented by `CometNativeScanExec` and the contrib's `CometDeltaNativeScanExec` -- without
931+
* it, [[PlanDataInjector.findAllPlanData]] cannot collect the per-partition tasks and the
932+
* parent's native execution receives an empty input. (`CometIcebergNativeScanExec` does NOT use
933+
* this trait; it has a dedicated `findAllPlanData` case.)
934+
*
935+
* Each implementation also resolves its own DPP subqueries via `ensureSubqueriesResolved`
936+
* (overridden from [[CometLeafExec]]) before `commonData`/`perPartitionData` are read.
937+
*/
938+
trait CometScanWithPlanData {
939+
def sourceKey: String
940+
def commonData: Array[Byte]
941+
def perPartitionData: Array[Array[Byte]]
942+
943+
// DPP / partition filters that may carry AQE SubqueryAdaptiveBroadcast
944+
// subqueries needing rewrite by CometPlanAdaptiveDynamicPruningFilters.
945+
// Default empty: scans with dedicated handling (CometNativeScanExec,
946+
// CometIcebergNativeScanExec) don't use this path.
947+
def dynamicPruningFilters: Seq[Expression] = Nil
948+
949+
// Install rewritten DPP filters on this scan. Implementers whose filters live
950+
// in a @transient field (which TreeNode.makeCopy can't carry, #3510) update
951+
// them via a transient side-channel and return `this` -- so the optimizer
952+
// rule's rewrite lands on the SAME instance that executes, instead of a copy
953+
// that gets dropped when the enclosing native block is rebuilt. Only called
954+
// when `dynamicPruningFilters` is non-empty, so the default is never reached
955+
// for scans that leave it empty.
956+
def withDynamicPruningFilters(filters: Seq[Expression]): SparkPlan =
957+
throw new UnsupportedOperationException(
958+
s"${getClass.getSimpleName} exposes dynamicPruningFilters but does not " +
959+
"override withDynamicPruningFilters")
960+
}
961+
884962
abstract class CometUnaryExec extends CometNativeExec with UnaryExecNode
885963

886964
abstract class CometBinaryExec extends CometNativeExec with BinaryExecNode
Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.spark.sql.comet
21+
22+
import org.scalatest.funsuite.AnyFunSuite
23+
24+
import org.apache.spark.sql.catalyst.expressions.Expression
25+
26+
import org.apache.comet.serde.OperatorOuterClass.Operator
27+
28+
/**
29+
* Unit coverage for the core SPI that lets out-of-tree contrib leaf scans (e.g. the Delta
30+
* contrib's `CometDeltaNativeScanExec`) participate in Comet native planning without core holding
31+
* a compile-time reference to them: the [[CometScanWithPlanData]] trait contract and the
32+
* reflective `DeltaPlanDataInjector` slot in [[PlanDataInjector]]'s registry.
33+
*
34+
* Deliberately does not exercise the per-op injector registry mechanics; that surface is owned by
35+
* `PlanDataInjectorSuite`.
36+
*/
37+
class CometScanWithPlanDataSuite extends AnyFunSuite {
38+
39+
/** Minimal trait implementation that opts into none of the optional DPP behaviour. */
40+
private class StubScan extends CometScanWithPlanData {
41+
override def sourceKey: String = "stub-key"
42+
override def commonData: Array[Byte] = Array.emptyByteArray
43+
override def perPartitionData: Array[Array[Byte]] = Array.empty
44+
}
45+
46+
test(
47+
"CometScanWithPlanData defaults: no DPP filters and withDynamicPruningFilters is unsupported") {
48+
val scan = new StubScan
49+
assert(scan.dynamicPruningFilters == Nil, "a scan must expose no DPP filters by default")
50+
// The default exists only so the DPP rule never calls it for scans that leave
51+
// dynamicPruningFilters empty; if a scan does expose filters it must override this.
52+
val ex = intercept[UnsupportedOperationException] {
53+
scan.withDynamicPruningFilters(Seq.empty[Expression])
54+
}
55+
assert(
56+
ex.getMessage.contains("withDynamicPruningFilters"),
57+
s"default failure must name the un-overridden method, got: ${ex.getMessage}")
58+
}
59+
60+
test("contrib DeltaPlanDataInjector slot is absent in default builds and declines cleanly") {
61+
// PlanDataInjector appends a reflective `DeltaPlanDataInjector$` slot to its registry only
62+
// when the contrib-delta build bundled it. A default build must not contain the class, and
63+
// the reflective lookup must degrade to a plain ClassNotFoundException (-> None) so the
64+
// registry keeps only its built-in injectors with no error.
65+
val deltaSlotAbsent =
66+
try {
67+
// scalastyle:off classforname
68+
Class.forName("org.apache.spark.sql.comet.DeltaPlanDataInjector$")
69+
// scalastyle:on classforname
70+
false
71+
} catch {
72+
case _: ClassNotFoundException => true
73+
}
74+
assert(
75+
deltaSlotAbsent,
76+
"default (non -Pcontrib-delta) build must not carry the contrib DeltaPlanDataInjector")
77+
78+
// With no contrib injector present, a non-scan operator tree must pass through untouched --
79+
// i.e. the reflective slot contributes zero surface to the default planner.
80+
val child = Operator.newBuilder().setPlanId(2).build()
81+
val root = Operator.newBuilder().setPlanId(1).addChildren(child).build()
82+
assert(
83+
PlanDataInjector.injectPlanData(root, Map.empty, Map.empty) == root,
84+
"non-scan operator tree must be returned unchanged with no contrib injector present")
85+
}
86+
}

0 commit comments

Comments
 (0)