andygrove commented on code in PR #6071: URL: https://github.com/apache/datafusion-comet/pull/6071#discussion_r4073736402
########## 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: `CometExecRDD.compute` injects per-partition planning data for native scans (`PlanDataInjector.injectPlanData`), so every task in those stages sends different bytes. With the cache on, each of those tasks misses, copies its full plan into a new entry, and evicts entries that other stages could have reused. Could we skip the cache when per-partition plan data was injected, or key on the base plan instead? A test showing that a stage with distinct per-task bytes doesn't evict a reusable entry would help guard this. ########## 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: If the shared build fails, for example when `convert_tree` hits `Unexpected operator in shared tree`, the `?` here fails the task. Could we fall back to `planner.create_plan` in that case and log it? Admission and the planner agree today, but a later planner change that adds a wrapper operator would otherwise break queries whenever the flag is on. ########## 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: Could you explain the change from the limit plan to a bare scan plan here? It looks like the original limit case is no longer covered by this lifecycle test. ########## 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: These tests use `checkSparkAnswerAndOperator`, which passes whether or not a shared tree was used. I think the post-shuffle Final aggregates and joins here read through `ShuffleScan`, so they probably run on private plans even with the flag on. Could we expose a small counter or test hook for shared-tree binds and assert on it? Asserting the expected fallback cases would be useful too. ########## 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: This override (and the matching one in `CometExecIteratorLifecycleSuite`) turns on both flags for every existing test in the suite. The flag-off path, which is what users run by default, loses that coverage here. Could we enable the flags only in the new tests, or run the affected tests in both modes? -- 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]
