viirya commented on code in PR #6071:
URL: https://github.com/apache/datafusion-comet/pull/6071#discussion_r4075324481
##########
native/core/src/execution/jni_api.rs:
##########
@@ -967,11 +994,35 @@ pub unsafe extern "system" fn
Java_org_apache_comet_Native_executePlan(
.with_shuffle_partition_pusher(
exec_context.shuffle_partition_pusher.clone(),
);
- let (scans, shuffle_scans, root_op) = planner.create_plan(
- &exec_context.spark_plan,
- &mut exec_context.input_sources.clone(),
- exec_context.partition_count,
- )?;
+ let (scans, shuffle_scans, root_op) =
+ if let Some(key) = &exec_context.shared_plan_key {
+ let shared = super::shared_pipeline::get_or_build(
Review Comment:
Implemented in ffdbbd2df833c99902c44dbc1f436520342badf3. Shared-tree
construction/conversion errors now log a warning and fall back to private
planning. The regression test introduces an unsupported wrapper around a real
DataFusion tree, verifies that conversion fails without leaving a registry
entry, and verifies the private plan's output. Binding and execution errors
still propagate: once task-owned input streams may have been imported, retrying
them would not be safe.
##########
native/core/src/execution/plan_cache.rs:
##########
@@ -0,0 +1,455 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Executor-local reuse of immutable protobuf plans and bounded single-flight
cache (#1204).
+//!
+//! The definition cache here owns protobuf data only. `shared_pipeline`
separately uses the
+//! generic cache for audited immutable physical trees. Neither cache may
retain task resources.
+//! Different partitions' scan payloads must remain different definition-cache
keys.
+
+use std::collections::HashMap;
+use std::sync::{Arc, LazyLock};
+use std::time::Instant;
+
+use datafusion_comet_proto::spark_operator::Operator;
+use once_cell::sync::OnceCell;
+use parking_lot::Mutex;
+
+use super::operators::ExecutionError;
+use super::serde::deserialize_op;
+
+const MAX_ENTRIES: usize = 64;
+const MAX_ENCODED_BYTES: usize = 8 * 1024 * 1024;
+
+// Process-local because tasks do not share a SessionContext. This cache owns
only immutable
+// protobuf data, never storage clients, JNI references or task resources.
Exact bytes are
+// compared, so neither hash collisions nor another query's configuration can
change the decoded
+// result. Retention is bounded by entry count and encoded bytes, and
release_runtime clears it.
+// The byte budget accounts for keys, not the decoded Rust heap (which can be
larger).
+static PLAN_CACHE: LazyLock<PlanCache<Operator>> =
+ LazyLock::new(|| PlanCache::new(MAX_ENTRIES, MAX_ENCODED_BYTES));
+
+pub(super) fn decode_plan(
Review Comment:
Addressed in ffdbbd2df833c99902c44dbc1f436520342badf3 by removing the
decoded-plan cache and its config entirely. The standalone measurements showed
only approximately 0.2%–1.4% lower elapsed time than both flags disabled, which
does not justify maintaining this cache. Injected per-task plan bytes are now
decoded without entering a native definition cache, so they cannot evict
reusable definitions. The physical registry still rejects native file scans.
##########
spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala:
##########
@@ -90,6 +93,114 @@ class CometExecSuite extends CometTestBase {
}
}
+ test("native plan cache setting crosses JNI for both enabled and disabled
execution") {
+ for (enabled <- Seq("true", "false")) {
+ withSQLConf(CometConf.COMET_EXEC_PLAN_CACHE_ENABLED.key -> enabled) {
+ val configs =
ConfigMap.parseFrom(CometExecIterator.serializeCometSQLConfs())
+
assert(configs.getEntriesMap.get(CometConf.COMET_EXEC_PLAN_CACHE_ENABLED.key)
== enabled)
+ withParquetTable((0 until 32).map(i => (i, i + 1)),
"plan_cache_input") {
+ checkSparkAnswerAndOperator(
+ sql("SELECT _1 + 1 FROM plan_cache_input WHERE _2 > 8"),
+ Seq(classOf[CometProjectExec]))
+ }
+ }
+ }
+ }
+
+ test("shared native pipelines across task waves and AQE") {
+ for (enabled <- Seq("true", "false"); aqe <- Seq("true", "false")) {
+ withSQLConf(
+ CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key -> enabled,
+ CometConf.COMET_EXEC_PLAN_CACHE_ENABLED.key -> "false",
+ CometConf.COMET_EXPLAIN_NATIVE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe,
+ SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "17") {
+ val configs =
ConfigMap.parseFrom(CometExecIterator.serializeCometSQLConfs())
+
assert(configs.getEntriesMap.get(CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key)
== enabled)
+ for (_ <- 0 until 2) {
+ val df = spark.range(0, 1000, 1, 16).where("id >
100").selectExpr("id + 10 AS value")
+ val (_, nativePlan) = checkSparkAnswerAndOperator(df,
Seq(classOf[CometProjectExec]))
+ val projects = stripAQEPlan(nativePlan).collect { case p:
CometProjectExec => p }
+ assert(projects.nonEmpty)
+ assert(projects.head.metrics("output_rows").value == 899L)
+ }
+ val empty = spark.range(0, 100, 1, 16).where("id < 0").selectExpr("id
+ 10 AS value")
+ checkSparkAnswerAndOperator(empty, Seq(classOf[CometProjectExec]))
+ // Stateful expressions in an otherwise eligible JVM-input block use
private plans.
+ val stateful = spark
+ .range(0, 100, 1, 16)
+ .selectExpr("spark_partition_id() AS partition",
"monotonically_increasing_id() AS id")
+ checkSparkAnswerAndOperator(stateful, Seq(classOf[CometProjectExec]))
+ }
+ }
+ }
+
+ test("shared DataFusion stateful operators across Spark partitions") {
Review Comment:
Added `shared_plan_tasks` SQL metrics and explicit path assertions in
ffdbbd2df833c99902c44dbc1f436520342badf3. Projection/filter tests assert the
exact task count; positive sort, post-shuffle Final aggregate, and broadcast
hash join tests arrange eligible JVM inputs and assert nonzero counts only when
sharing is enabled, with AQE both on and off. A separate test verifies that the
executed block actually contains a serialized ShuffleScan and asserts zero
shared tasks. Native-scan and stateful-expression fallback cases also assert
zero.
This metric counts tasks bound to a shared tree, not registry hits; native
tests continue to verify actual operator identity. The broader
output-equivalence tests have been renamed so their names do not imply that
every case was admitted. I also found that the direct-read setting alone is not
enough to infer the input representation: some exchange plans still serialize
as ordinary Scan inputs, so the fallback test checks the serialized plan itself.
##########
spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala:
##########
@@ -63,7 +63,10 @@ class CometExecSuite extends CometTestBase {
override protected def test(testName: String, testTags: Tag*)(testFun: =>
Any)(implicit
pos: Position): Unit = {
super.test(testName, testTags: _*) {
- withSQLConf(CometConf.COMET_SHUFFLE_ENABLED.key -> "true") {
+ withSQLConf(
+ CometConf.COMET_SHUFFLE_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_PLAN_CACHE_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key -> "true") {
Review Comment:
Fixed in ffdbbd2df833c99902c44dbc1f436520342badf3. The suite-wide
sharing/cache overrides are removed, restoring default-path coverage for
existing tests. The cache flag itself is gone. New sharing tests enable the
feature locally, and the lifecycle regressions run with sharing disabled and
enabled. The lifecycle test's synthetic TaskContext now carries SQLConf through
local properties, as real Spark tasks do; otherwise its empty properties
silently selected the default setting.
##########
spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala:
##########
@@ -192,12 +235,15 @@ class CometExecIteratorLifecycleSuite extends
CometTestBase {
schema,
CometArrowStream.NATIVE_TIMEZONE,
"lifecycle-test")
- val limitOp =
- CometExecUtils.getLimitNativePlan(Seq(PrettyAttribute("test",
LongType)), 100).get
+ val scanOp =
Review Comment:
You are right: I replaced the limit with a scan to enter an admitted sharing
path, but that should have been an additional case. In
ffdbbd2df833c99902c44dbc1f436520342badf3, the limit case is restored and tested
with sharing off/on. A separate scan case is also tested in both modes, and
inspects the final native metrics report to assert whether a shared tree was
actually bound before the injected metrics-update failure. All nine lifecycle
tests pass.
--
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]
