rich7420 commented on code in PR #7544: URL: https://github.com/apache/datafusion-comet/pull/7544#discussion_r4528021374
########## spark/src/main/scala/org/apache/spark/sql/comet/CometNativeExecutorsExec.scala: ########## @@ -0,0 +1,657 @@ +/* + * 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. + */ + +package org.apache.spark.sql.comet + +import java.util.IdentityHashMap + +import scala.collection.mutable.ArrayBuffer +import scala.jdk.CollectionConverters._ + +import org.apache.spark.{Partition, SparkContext, TaskContext} +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, DynamicPruning, DynamicPruningExpression, Expression, Literal} +import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, RangePartitioning, SinglePartition} +import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec} +import org.apache.spark.sql.execution.{ColumnarToRowExec, CommandResultExec, DeserializeToObjectExec, LeafExecNode, ScalarSubquery, SparkPlan} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, QueryStageExec} +import org.apache.spark.sql.execution.command.{DataWritingCommandExec, ExecutedCommandExec} +import org.apache.spark.sql.execution.datasources.v2.V2CommandExec +import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.StructField +import org.apache.spark.sql.vectorized.ColumnarBatch + +import com.google.protobuf.ByteString +import com.google.protobuf.Descriptors.FieldDescriptor +import com.google.protobuf.Message + +import org.apache.comet.{CometConf, CometExecIterator, CometExplainInfo, CometRuntimeException, NativeExecutorsClient} +import org.apache.comet.serde.{ExprOuterClass, PartitioningOuterClass, QueryPlanSerde} +import org.apache.comet.serde.Config.ConfigMap +import org.apache.comet.serde.NativeExecutors.{QueryRequest, ResultEndpoint, ResultEndpoints, Stage, StageInput} +import org.apache.comet.serde.OperatorOuterClass +import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.operator.schema2Proto +import org.apache.comet.vector.NativeUtil + +/** + * Runs a whole query on a Comet Native Executor cluster instead of on Spark's executors. See + * https://github.com/apache/datafusion-comet/issues/7288. + * + * This leaf replaces the plan Comet would otherwise execute. Planning only checks that the plan + * can run on a cluster ([[CometNativeExecutors.validate]]). `doExecuteColumnar` builds the query, + * which lists the files the scans read, submits it to the cluster's scheduler and waits for it on + * the driver; each partition of the RDD it returns then reads one result partition from the + * executor that holds it. + * + * @param nativePlan + * the plan Comet would execute, which [[CometNativeExecutors.buildQuery]] turns into the + * `QueryRequest` the cluster runs + * @param subqueryResults + * the query's scalar subqueries. They are expressions of this node, so Spark runs them before + * it executes the node, and their results replace them in the query before it is submitted. + */ +case class CometNativeExecutorsExec( + override val output: Seq[Attribute], + nativePlan: SparkPlan, + schedulerUrl: String, + subqueryResults: Seq[Expression] = Nil) + extends LeafExecNode { + + override def supportsColumnar: Boolean = true + + override protected def doExecute(): RDD[InternalRow] = + throw new UnsupportedOperationException(s"$nodeName only produces columnar output") + + override protected def doExecuteColumnar(): RDD[ColumnarBatch] = { + NativeExecutorsClient.requireNativeSupport() + val request = CometNativeExecutors.withSubqueryResults( + CometNativeExecutors.buildQuery(nativePlan).toByteArray, Review Comment: A query with a subquery is encoded here, parsed in `withSubqueryResults`, and encoded again. Could `withSubqueryResults` take and return the `QueryRequest`, with a single `toByteArray` call before `submit`? ########## spark/src/main/scala/org/apache/spark/sql/comet/CometNativeExecutorsExec.scala: ########## @@ -0,0 +1,657 @@ +/* + * 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. + */ + +package org.apache.spark.sql.comet + +import java.util.IdentityHashMap + +import scala.collection.mutable.ArrayBuffer +import scala.jdk.CollectionConverters._ + +import org.apache.spark.{Partition, SparkContext, TaskContext} +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, DynamicPruning, DynamicPruningExpression, Expression, Literal} +import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, RangePartitioning, SinglePartition} +import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec} +import org.apache.spark.sql.execution.{ColumnarToRowExec, CommandResultExec, DeserializeToObjectExec, LeafExecNode, ScalarSubquery, SparkPlan} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, QueryStageExec} +import org.apache.spark.sql.execution.command.{DataWritingCommandExec, ExecutedCommandExec} +import org.apache.spark.sql.execution.datasources.v2.V2CommandExec +import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.StructField +import org.apache.spark.sql.vectorized.ColumnarBatch + +import com.google.protobuf.ByteString +import com.google.protobuf.Descriptors.FieldDescriptor +import com.google.protobuf.Message + +import org.apache.comet.{CometConf, CometExecIterator, CometExplainInfo, CometRuntimeException, NativeExecutorsClient} +import org.apache.comet.serde.{ExprOuterClass, PartitioningOuterClass, QueryPlanSerde} +import org.apache.comet.serde.Config.ConfigMap +import org.apache.comet.serde.NativeExecutors.{QueryRequest, ResultEndpoint, ResultEndpoints, Stage, StageInput} +import org.apache.comet.serde.OperatorOuterClass +import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.operator.schema2Proto +import org.apache.comet.vector.NativeUtil + +/** + * Runs a whole query on a Comet Native Executor cluster instead of on Spark's executors. See + * https://github.com/apache/datafusion-comet/issues/7288. + * + * This leaf replaces the plan Comet would otherwise execute. Planning only checks that the plan + * can run on a cluster ([[CometNativeExecutors.validate]]). `doExecuteColumnar` builds the query, + * which lists the files the scans read, submits it to the cluster's scheduler and waits for it on + * the driver; each partition of the RDD it returns then reads one result partition from the + * executor that holds it. + * + * @param nativePlan + * the plan Comet would execute, which [[CometNativeExecutors.buildQuery]] turns into the + * `QueryRequest` the cluster runs + * @param subqueryResults + * the query's scalar subqueries. They are expressions of this node, so Spark runs them before + * it executes the node, and their results replace them in the query before it is submitted. + */ +case class CometNativeExecutorsExec( + override val output: Seq[Attribute], + nativePlan: SparkPlan, + schedulerUrl: String, + subqueryResults: Seq[Expression] = Nil) + extends LeafExecNode { + + override def supportsColumnar: Boolean = true + + override protected def doExecute(): RDD[InternalRow] = + throw new UnsupportedOperationException(s"$nodeName only produces columnar output") + + override protected def doExecuteColumnar(): RDD[ColumnarBatch] = { + NativeExecutorsClient.requireNativeSupport() + val request = CometNativeExecutors.withSubqueryResults( + CometNativeExecutors.buildQuery(nativePlan).toByteArray, + subqueryResults.collect { case sub: ScalarSubquery => sub }) + val endpoints = + ResultEndpoints.parseFrom(new NativeExecutorsClient().submit(schedulerUrl, request)) + new CometNativeExecutorsRDD( + sparkContext, + endpoints.getEndpointsList.asScala.toSeq, + output.length) + } + + override def stringArgs: Iterator[Any] = + Iterator(output, schedulerUrl) +} + +private class CometNativeExecutorsPartition(override val index: Int, val endpoint: ResultEndpoint) + extends Partition + +/** + * The result partitions of a query that ran on a Comet Native Executor cluster, each read from + * the executor that holds it. + */ +private class CometNativeExecutorsRDD( + sc: SparkContext, + endpoints: Seq[ResultEndpoint], + numColumns: Int) + extends RDD[ColumnarBatch](sc, Nil) { + + override protected def getPartitions: Array[Partition] = + endpoints.zipWithIndex.map { case (endpoint, index) => + new CometNativeExecutorsPartition(index, endpoint): Partition + }.toArray + + override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = { + val partition = split.asInstanceOf[CometNativeExecutorsPartition] + // A Spark executor needs a libcomet built with the feature too. + NativeExecutorsClient.requireNativeSupport() + val client = new NativeExecutorsClient() + val nativeUtil = new NativeUtil() + // The task's completion listener closes nativeUtil, but only once there is one to register. + val handle = + try { + client.openPartition( + partition.endpoint.getLocation, + partition.endpoint.getTicket.toByteArray) + } catch { + case t: Throwable => + nativeUtil.close() + throw t + } + val batches = new Iterator[ColumnarBatch] with AutoCloseable { + private var current: Option[ColumnarBatch] = None + private var pending: Option[ColumnarBatch] = None + private var closed = false + + override def hasNext: Boolean = { + if (pending.isEmpty && !closed) { + // The consumer has moved past the batch it last got. + current.foreach(_.close()) + current = None + pending = nativeUtil.getNextBatch( + numColumns, + (arrays, schemas) => client.nextBatch(handle, arrays, schemas)) + if (pending.isEmpty) { + close() + } + } + pending.isDefined + } + + override def next(): ColumnarBatch = { + if (!hasNext) { + throw new NoSuchElementException("No more batches in this result partition") + } + current = pending + pending = None + current.get + } + + override def close(): Unit = if (!closed) { + closed = true + try { + current.foreach(_.close()) + current = None + pending.foreach(_.close()) + pending = None + } finally { + // Each step runs even if an earlier one throws. + try client.closePartition(handle) + finally nativeUtil.close() + } + } + } + context.addTaskCompletionListener[Unit](_ => batches.close()) + batches + } +} + +/** + * Plans queries for Comet Native Executors: turns the plan Comet would execute into the stage + * graph a cluster runs, and refuses any query that cannot run entirely in the Comet pipeline. + */ +object CometNativeExecutors { + + /** + * Messages of native plans that call back into the JVM, which a cluster doesn't have. A scalar + * subquery's `Subquery` is not one of them: the driver runs the subquery and puts its result in + * the plan before submitting it. The native planner could refuse them all instead + * (https://github.com/apache/datafusion-comet/issues/7375). + */ + private val jvmCallbacks = Set("JvmScalarUdf", "ArrowPythonUdf") + + /** + * Whether Comet runs `plan` on a Comet Native Executor cluster. It must be a whole query: not a + * command, not a stage AQE has split off, and not a plan already replaced. The caller checks + * that Comet executes the plan at all. + */ + def applies(conf: SQLConf, plan: SparkPlan, queryStagePrep: Boolean): Boolean = + CometConf.COMET_NATIVE_EXECUTORS_ENABLED.get(conf) && + !(conf.adaptiveExecutionEnabled && !queryStagePrep && plan.isInstanceOf[Exchange]) && + !plan.exists { + case _: CometNativeExecutorsExec | _: AdaptiveSparkPlanExec | _: QueryStageExec | + _: CommandResultExec | _: ExecutedCommandExec | _: DataWritingCommandExec | + _: V2CommandExec => + true + case _ => false + } + + /** + * Replaces `converted`, the plan Comet would execute for `original`, with a leaf that runs it + * on a Comet Native Executor cluster. An operator on top that only turns the query's result + * into JVM objects, as `Dataset.rdd` adds, stays on the driver above the leaf. + */ + def offload(original: SparkPlan, converted: SparkPlan): SparkPlan = { + NativeExecutorsClient.requireNativeSupport() + offloadPlan(original, converted) + } + + private def offloadPlan(original: SparkPlan, converted: SparkPlan): SparkPlan = + converted match { + case toObjects: DeserializeToObjectExec => + toObjects.withNewChildren(Seq(offloadPlan(toObjects.child, toObjects.child))) + case native => + validate(native) + val exec = CometNativeExecutorsExec( + native.output, + native, + CometConf.COMET_NATIVE_EXECUTORS_SCHEDULER_URL.get(), + scalarSubqueries(native)) + native.logicalLink.orElse(original.logicalLink).foreach(exec.setLogicalLink) + exec + } + + /** + * Checks that a cluster can run `plan`, the plan Comet would execute, without running anything: + * it lists no files and resolves no subquery, so it is safe at planning time. + * + * @throws CometRuntimeException + * naming the first operator that cannot run on a cluster + */ + def validate(plan: SparkPlan): Unit = stageGraph(plan, withFiles = false) Review Comment: The Strict Scala warnings (Spark 3.5) job fails here with "discarded non-Unit value". Could `validate` call `stageGraph` and then return `()`? That change also clears the four `@nowarn` errors. ########## spark/src/main/scala/org/apache/comet/rules/CometRule.scala: ########## @@ -52,16 +53,28 @@ object CometRule { CometCacheColumnarRule(preview)) /** - * Whether plan-only mode applies to `plan`, so that Comet only reports the plan it would - * execute and Spark executes `plan` unchanged. Mirrors the conversion rules' own guards; - * plan-only is scoped to exec being enabled. + * Whether the conversion rules convert `plan`: they mirror these guards. Plan-only mode and + * Comet Native Executors, which each rely on the conversion, are scoped to it. */ - private[comet] def planOnlyApplies(conf: SQLConf, plan: SparkPlan): Boolean = - CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(conf) && - isCometLoaded(conf) && + private def conversionApplies(conf: SQLConf, plan: SparkPlan): Boolean = + isCometLoaded(conf) && !ShimCometStreaming.isStreamingPlan(plan) && CometConf.COMET_EXEC_ENABLED.get(conf) + /** + * Whether plan-only mode applies to `plan`, so that Comet only reports the plan it would + * execute and Spark executes `plan` unchanged. + */ + private[comet] def planOnlyApplies(conf: SQLConf, plan: SparkPlan): Boolean = + CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(conf) && conversionApplies(conf, plan) + + /** Whether `plan` runs on a Comet Native Executor cluster instead of Spark's executors. */ + private def nativeExecutorsApply( Review Comment: This guard looks like the fix #7557 asks for, and `CometNativeExecutorsPlanSuite` tests it. Could the PR description close #7557 as well? ########## docs/source/contributor-guide/native_executors.md: ########## @@ -0,0 +1,131 @@ +<!-- +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. +--> + +# Comet Native Executors + +**Experimental.** Comet Native Executors run whole Spark queries without a JVM in the data plane. A +Spark driver plans each query with Catalyst and Comet, then submits the plan as a graph of stages +to a cluster of one `comet-scheduler` and any number of `comet-executor` processes, built on +[Apache DataFusion Ballista](https://github.com/apache/datafusion-ballista). Each stage is a +_block_: a run of Comet operators that Comet serializes as one Rust plan, with the exchanges between +blocks as the edges of the graph. The design is in +[issue #7288](https://github.com/apache/datafusion-comet/issues/7288), and +[`native/native-executors/README.md`](https://github.com/apache/datafusion-comet/blob/main/native/native-executors/README.md) +describes the cluster, how to build it and how to test it. + +This page is for contributors. It covers the Spark driver side: how to run a query on a cluster and +what does not work yet. The driver never links Ballista. It sends the stage graph to the +scheduler's Arrow Flight service through a small client in libcomet, and the Spark tasks that read +the results use the same client. The cluster features that some queries rely on are in +[#7289](https://github.com/apache/datafusion-comet/pull/7289) and are not part of the cluster in +this repository yet. [Current limitations](#current-limitations) says which. + +## Running it + +1. Build libcomet with the `native-executors` Cargo feature. It is off by default, so the standard + build links no gRPC client. A libcomet without it fails the first query with an error that names + the feature. For example, `make COMET_FEATURES=native-executors` builds libcomet and the Comet + jar that bundles it. The Spark executors that read the results need this libcomet too. +2. Build and start the cluster. The `comet-scheduler` and `comet-executor` binaries are their own + Cargo workspace, so the main Comet build does not depend on Ballista. The cluster and the driver + must come from the same Comet build, which the scheduler checks. See + [`native/native-executors/README.md`](https://github.com/apache/datafusion-comet/blob/main/native/native-executors/README.md) + for how to build the binaries and for the environment variables that configure them. +3. Point Spark at the scheduler's Flight service and turn the feature on: + + ```shell + --conf spark.comet.nativeExecutors.enabled=true \ + --conf spark.comet.nativeExecutors.scheduler.url=http://scheduler-host:50060 + ``` + + Both configs are internal while the feature is experimental. Comet itself must be enabled as + usual, with native shuffle (`spark.comet.shuffle.mode=native`, as in the + [installation guide](../user-guide/latest/installation.md)), because the driver turns only + native shuffle exchanges into the edges of the stage graph. + +4. Check that the query ran on the cluster: its physical plan has a `CometNativeExecutors` leaf + (the `CometNativeExecutorsExec` class) where the Comet plan would otherwise be. A query that + cannot run entirely in the Comet pipeline and on the cluster fails instead, with a message that + names the operator and Comet's reasons for it. + +`CometNativeExecutorsSuite` runs queries against a cluster it starts from the binaries in +`COMET_NATIVE_EXECUTORS_BIN_DIR`. Its header comment gives the commands. It is not part of CI. Its +tests that need a cluster are skipped when the variable is not set. The tests that check what the +driver refuses to plan run without a cluster, but they still need a libcomet built with the +`native-executors` feature. Its TPC-H tests need the cluster of +[#7289](https://github.com/apache/datafusion-comet/pull/7289): on the cluster in this repository, +only q1, q6, q11 and q15 pass. `CometNativeExecutorsPlanSuite` runs in CI. It takes the plan Comet +would execute and checks the stage graph, the queries the driver refuses, and how a scalar +subquery's result goes into the plan, with neither a cluster nor the feature. + +## Current limitations + +Comet Native Executors are a first version. The sub-issues of +[#7288](https://github.com/apache/datafusion-comet/issues/7288) track what is missing: + +- Whole queries only. See [below](#whole-queries-only). +- Only native shuffle exchanges. A query whose exchanges use JVM shuffle is refused. +- `LIMIT` and `show()`: [#7365](https://github.com/apache/datafusion-comet/issues/7365). +- Global sorts. The cluster in this repository reads a range exchange as one partition, so every + global sort runs in a single task. + [#7289](https://github.com/apache/datafusion-comet/pull/7289) adds a cluster that sorts the + partitions in parallel and merges them, except for sorts by floats or nested types, which still + run in a single task: [#7366](https://github.com/apache/datafusion-comet/issues/7366). +- Hash partitioning is DataFusion's, not Spark-compatible: + [#7368](https://github.com/apache/datafusion-comet/issues/7368). +- A small broadcast side runs in every task that reads it: + [#7369](https://github.com/apache/datafusion-comet/issues/7369). Moving a join whose build side + turns out small into the stage that produces its probe side is in + [#7289](https://github.com/apache/datafusion-comet/pull/7289). +- Dynamic partition pruning: [#7370](https://github.com/apache/datafusion-comet/issues/7370). A scan + with a dynamic pruning filter is refused. +- Writes: [#7371](https://github.com/apache/datafusion-comet/issues/7371). Commands and writes stay + in Spark. +- Cancellation: [#7372](https://github.com/apache/datafusion-comet/issues/7372). The driver thread + waits for the cluster in native code until the query finishes, so cancelling the Spark job or + interrupting the thread does not stop the wait, and a killed Spark job leaves its query running + on the cluster. +- Errors. A failure on the cluster reaches Spark as a `CometNativeException` with the cluster's Review Comment: Since rebuilding Spark's exception is planned as a follow-up, could this item link to a sub-issue of #7288, like the cancellation and metrics items? I couldn't find one for errors. -- 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]
