rich7420 commented on code in PR #7544: URL: https://github.com/apache/datafusion-comet/pull/7544#discussion_r4521540891
########## spark/src/test/scala/org/apache/comet/exec/CometNativeExecutorsSuite.scala: ########## @@ -0,0 +1,451 @@ +/* + * 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.comet.exec + +import java.io.File +import java.net.{ServerSocket, Socket} +import java.nio.file.Files +import java.util.concurrent.TimeUnit + +import scala.util.Try + +import org.apache.spark.SparkConf +import org.apache.spark.sql.{CometTestBase, DataFrame, Row} +import org.apache.spark.sql.comet.CometNativeExecutorsExec +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.functions.{col, count, max, sum} +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.CometConf + +/** + * Runs queries on a Comet Native Executor cluster (see + * https://github.com/apache/datafusion-comet/issues/7288): a comet-scheduler and a comet-executor + * started from the binaries in the directory `COMET_NATIVE_EXECUTORS_BIN_DIR` names. The tests + * are skipped when it isn't set. To run them: + * + * {{{ + * cargo build --manifest-path native/Cargo.toml -p datafusion-comet --features native-executors + * CARGO_TARGET_DIR=native/target cargo build --manifest-path native/native-executors/Cargo.toml + * COMET_NATIVE_EXECUTORS_BIN_DIR=native/target/debug ./mvnw test -Dtest=none \ + * -Dsuites=org.apache.comet.exec.CometNativeExecutorsSuite + * }}} + */ +class CometNativeExecutorsSuite extends CometTestBase with AdaptiveSparkPlanHelper { + + private val binDir: Option[File] = + sys.env.get("COMET_NATIVE_EXECUTORS_BIN_DIR").map(new File(_).getAbsoluteFile) + private var processes: Seq[Process] = Nil + private var schedulerUrl: String = _ + + override protected def sparkConf: SparkConf = { + val conf = super.sparkConf + conf.set(CometConf.COMET_SHUFFLE_MODE.key, "native") + conf.set(SQLConf.SHUFFLE_PARTITIONS.key, "4") + conf + } + + override def beforeAll(): Unit = { + super.beforeAll() + binDir.foreach(startCluster) + } + + override def afterAll(): Unit = { + try { + // The executor first, so that it doesn't report to a scheduler that has gone. + processes.reverse.foreach { process => + process.destroy() + if (!process.waitFor(10, TimeUnit.SECONDS)) { + process.destroyForcibly() + } + } + } finally { + super.afterAll() + } + } + + private def startCluster(dir: File): Unit = { + val schedulerPort = freePort() + val flightPort = freePort() + val executorPort = freePort() + val executorGrpcPort = freePort() + val logs = Files.createTempDirectory("comet-native-executors").toFile + logInfo(s"Comet Native Executor cluster logs: $logs") + start( + new File(dir, "comet-scheduler"), + Map( + "COMET_SCHEDULER_BIND_HOST" -> "127.0.0.1", + "COMET_SCHEDULER_PORT" -> schedulerPort.toString, + "COMET_SCHEDULER_FLIGHT_PORT" -> flightPort.toString), + new File(logs, "scheduler.log")) + waitForPort(schedulerPort) + waitForPort(flightPort) + start( + new File(dir, "comet-executor"), + Map( + "COMET_EXECUTOR_BIND_HOST" -> "127.0.0.1", + "COMET_EXECUTOR_PORT" -> executorPort.toString, + "COMET_EXECUTOR_GRPC_PORT" -> executorGrpcPort.toString, + "COMET_SCHEDULER_HOST" -> "127.0.0.1", + "COMET_SCHEDULER_PORT" -> schedulerPort.toString, + "COMET_EXECUTOR_TASKS" -> "4", + "COMET_EXECUTOR_WORK_DIR" -> new File(logs, "work").getPath), + new File(logs, "executor.log")) + waitForPort(executorPort) + schedulerUrl = s"http://127.0.0.1:$flightPort"; + } + + private def start(binary: File, env: Map[String, String], log: File): Unit = { + val builder = new ProcessBuilder(binary.getPath) + env.foreach { case (key, value) => builder.environment().put(key, value) } + builder.redirectErrorStream(true).redirectOutput(log) + processes :+= builder.start() + } + + private def freePort(): Int = { + val socket = new ServerSocket(0) + try socket.getLocalPort + finally socket.close() + } + + private def waitForPort(port: Int): Unit = { + val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(60) + while (Try(new Socket("127.0.0.1", port).close()).isFailure) { + assert(System.nanoTime() < deadline, s"Nothing is listening on port $port") + Thread.sleep(100) + } + } + + private def nativeExecutorsTest(name: String)(body: => Unit): Unit = test(name) { + assume(binDir.isDefined, "COMET_NATIVE_EXECUTORS_BIN_DIR is not set") + body + } + + private def onCluster(body: => Unit): Unit = withSQLConf( + CometConf.COMET_NATIVE_EXECUTORS_ENABLED.key -> "true", + // Without a cluster, only tests that the planner refuses run, and those never connect. + CometConf.COMET_NATIVE_EXECUTORS_SCHEDULER_URL.key -> + Option(schedulerUrl).getOrElse("http://localhost:50060";))(body) + + /** + * Checks that the cluster runs `query` and returns what Spark returns, in order if `ordered`. + */ + private def checkOnCluster(query: => DataFrame, ordered: Boolean = false): Unit = + checkClusterReturns(query, query.collect().toSeq, ordered) Review Comment: The doc says this checks Spark's answer, but `query.collect()` runs with Comet enabled. A bug shared by Comet and the cluster could go undetected. Could we collect with `COMET_ENABLED=false`, as the TPC-H test does? ########## native/core/src/execution/native_executors/client.rs: ########## @@ -0,0 +1,271 @@ +// 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. + +//! The JNI client a Spark driver uses to run queries on a Comet Native Executor cluster: the +//! native methods of `org.apache.comet.NativeExecutorsClient`. +//! +//! The driver submits a query to the scheduler's Arrow Flight service, which answers with one +//! endpoint per result partition, and then reads each partition from its executor with `DoGet`. + +use std::time::Duration; + +use arrow_flight::decode::FlightRecordBatchStream; +use arrow_flight::error::FlightError; +use arrow_flight::flight_service_client::FlightServiceClient; +use arrow_flight::{FlightClient, FlightDescriptor, Ticket}; +use datafusion_comet_proto::spark_native_executors::{ + QueryRequest, ResultEndpoint, ResultEndpoints, +}; +use datafusion_comet_proto::PROTO_FINGERPRINT; +use futures::StreamExt; +use jni::objects::{JByteArray, JClass, JLongArray, JString, ReleaseMode}; +use jni::sys::{jbyteArray, jlong}; +use jni::EnvUnowned; +use prost::Message; +use tonic::transport::{Channel, Endpoint}; +use tonic::Code; + +use super::COMET_VERSION; +use crate::errors::{try_unwrap_or_throw, CometError, CometResult}; +use crate::execution::jni_api::{export_batch, get_runtime}; + +fn client_error(context: &str, error: impl std::fmt::Display) -> CometError { + CometError::Internal(format!("Comet Native Executors: {context}: {error}")) +} + +/// How long to wait for a connection to a scheduler or an executor. +const CONNECT_TIMEOUT: Duration = Duration::from_secs(30); + +/// How often to probe a connection that carries no traffic, at the TCP level and at the HTTP/2 +/// level. A query can run for a long time before the scheduler answers its submission, so a +/// request has no overall timeout. These probes are what notice a scheduler or an executor that +/// died or became unreachable without closing the connection. +const KEEP_ALIVE_INTERVAL: Duration = Duration::from_secs(30); + +/// How long to wait for an answer to an HTTP/2 probe before giving up on the connection. +const KEEP_ALIVE_TIMEOUT: Duration = Duration::from_secs(20); + +/// The error for a failed read of a result partition. A partition the cluster has deleted gets a +/// message that names the setting to change. The scheduler says that a partition it stores has +/// expired, and an executor fails to open the file of a partition it has deleted. +fn read_error(context: &str, error: FlightError) -> CometError { + let expired = match &error { + FlightError::Tonic(status) => match status.code() { + Code::NotFound => status.message().contains("expired"), + Code::Internal => { + let message = status.message(); + message.contains("Failed to open partition file") && message.contains("NotFound") + } + _ => false, + }, + _ => false, + }; + if expired { + client_error( + context, + format!( + "the result partition expired on the cluster; raise \ + COMET_SCHEDULER_RESULT_RETENTION_SECONDS ({error})" + ), + ) + } else { + client_error(context, error) + } +} + +/// Connects to the Flight service at `location`, a Flight location such as +/// `grpc+tcp://host:port` or an `http://` URL. +async fn connect(location: &str) -> CometResult<FlightClient> { + let uri = match location.strip_prefix("grpc+tcp://") { + Some(address) => format!("http://{address}";), + None => location.to_string(), + }; + let channel: Channel = Endpoint::from_shared(uri) + .map_err(|e| client_error(&format!("invalid location {location}"), e))? + .connect_timeout(CONNECT_TIMEOUT) + .tcp_keepalive(Some(KEEP_ALIVE_INTERVAL)) + .http2_keep_alive_interval(KEEP_ALIVE_INTERVAL) + .keep_alive_timeout(KEEP_ALIVE_TIMEOUT) + .keep_alive_while_idle(true) + .connect() + .await + .map_err(|e| client_error(&format!("cannot connect to {location}"), e))?; + // Query plans and result batches can be larger than gRPC's default 4 MiB message limit. + let inner = FlightServiceClient::new(channel) + .max_decoding_message_size(usize::MAX) + .max_encoding_message_size(usize::MAX); + Ok(FlightClient::new_from_inner(inner)) +} + +/// Submits `query`, a serialized `QueryRequest`, to the scheduler at `scheduler_url` and waits +/// for it to finish. Returns where its result partitions are. +async fn submit(scheduler_url: &str, mut query: Vec<u8>) -> CometResult<ResultEndpoints> { + // Stamp the query with this build's version and plan format. A field that appears again + // replaces the first, so appending them spares decoding and encoding the whole query. + QueryRequest { + comet_version: COMET_VERSION.to_string(), + proto_fingerprint: PROTO_FINGERPRINT.to_string(), + ..Default::default() + } + .encode(&mut query) + .map_err(|e| client_error("bad query", e))?; + let mut client = connect(scheduler_url).await?; + let info = client + .get_flight_info(FlightDescriptor::new_cmd(query)) + .await + .map_err(|e| client_error("the query failed", e))?; Review Comment: With ANSI on, `id div k` raises `SparkArithmeticException [DIVIDE_BY_ZERO]` in Spark. Ballista keeps only the Debug string, so here it surfaces as `CometNativeException`. Could the driver rebuild the exception through `SparkErrorConverter`, like `CometExecIterator`, or should Current limitations track this? ########## spark/src/main/scala/org/apache/spark/sql/comet/CometNativeExecutorsExec.scala: ########## @@ -0,0 +1,622 @@ +/* + * 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 scala.util.Try + +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.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. `doExecuteColumnar` submits the + * query 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 query + * the serialized `QueryRequest` built by [[CometNativeExecutors.buildQuery]] + * @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], + query: Array[Byte], + 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( + query, + 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, s"${query.length} bytes", 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. + */ + 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 => + val exec = CometNativeExecutorsExec( + native.output, + buildQuery(native).toByteArray, + CometConf.COMET_NATIVE_EXECUTORS_SCHEDULER_URL.get(), + scalarSubqueries(native)) + native.logicalLink.orElse(original.logicalLink).foreach(exec.setLogicalLink) + exec + } + + /** + * The query a Comet Native Executor cluster runs for `plan`, the plan Comet would execute: one + * stage per block of Comet operators that run as one Rust plan, with an input for each exchange + * that feeds the block. + * + * @throws CometRuntimeException + * naming the first operator that cannot run on a cluster + */ + def buildQuery(plan: SparkPlan): QueryRequest = { + val builder = new StageGraphBuilder(plan) + val root = plan match { + case ColumnarToRowExec(child) => child + case CometColumnarToRowExec(child) => child + case CometNativeColumnarToRowExec(child) => child + case other => other + } + builder.stageOf(root) + QueryRequest + .newBuilder() + .putAllConf(nativeConf()) + .addAllStages(builder.stages.asJava) + .build() + } + + /** The scalar subqueries of `plan`'s expressions, through exchanges, each once. */ + private def scalarSubqueries(plan: SparkPlan): Seq[Expression] = { + val found = ArrayBuffer.empty[ScalarSubquery] + def walk(node: SparkPlan): Unit = { + node.expressions.foreach(_.foreach { + case sub: ScalarSubquery if !found.exists(_.exprId == sub.exprId) => found += sub + case _ => + }) + node match { + case reused: ReusedExchangeExec => walk(reused.child) + case _ => node.children.foreach(walk) + } + } + walk(plan) + found.toSeq + } + + /** + * `query` with each scalar subquery's `Subquery` expression replaced by a literal of its + * result. The subqueries must have run. Sending the results alongside the plan would spare + * rewriting it (https://github.com/apache/datafusion-comet/issues/7374). + */ + def withSubqueryResults(query: Array[Byte], subqueries: Seq[ScalarSubquery]): Array[Byte] = { + if (subqueries.isEmpty) { + return query + } + val literals: Map[Long, ExprOuterClass.Expr] = subqueries.map { sub => + // The result is already in Catalyst's internal form. + val literal = Literal(sub.eval(null), sub.dataType) + val expr = QueryPlanSerde + .exprToProto(literal, Seq.empty) + .getOrElse(throw new CometRuntimeException( + s"Comet Native Executors cannot pass the result of $sub to the cluster")) + sub.exprId.id -> expr + }.toMap + val replace: PartialFunction[Message, Message] = { + case expr: ExprOuterClass.Expr if expr.hasSubquery => + literals.getOrElse( + expr.getSubquery.getId, + throw new CometRuntimeException( + s"Comet Native Executors found no result for subquery ${expr.getSubquery.getId}")) + } + val request = QueryRequest.parseFrom(query).toBuilder + (0 until request.getStagesCount).foreach { i => + val stage = request.getStages(i) + val plan = transformMessages(stage.getPlan)(replace) + // A partition's plan is the stage's with its files added, so it holds a subquery only + // when the stage's plan does. + if (plan ne stage.getPlan) { + val replaced = stage.toBuilder.setPlan(plan.asInstanceOf[Operator]) + (0 until stage.getPartitionPlansCount).foreach { p => + val partitionPlan = Operator.parseFrom(stage.getPartitionPlans(p)) + replaced.setPartitionPlans(p, transformMessages(partitionPlan)(replace).toByteString) + } + request.setStages(i, replaced) + } + } + request.build().toByteArray + } + + /** + * The messages directly inside `message`, each with how to put a replacement for it in a + * builder of `message`. + */ + private def childMessages( + message: Message): Iterator[(Message, (Message.Builder, Message) => Unit)] = + message.getAllFields.asScala.iterator + .filter { case (field, _) => field.getJavaType == FieldDescriptor.JavaType.MESSAGE } + .flatMap { case (field, value) => + if (field.isRepeated) { + value.asInstanceOf[java.util.List[Message]].asScala.iterator.zipWithIndex.map { + case (child, index) => + ( + child, + (builder: Message.Builder, m: Message) => { + builder.setRepeatedField(field, index, m) + () + }) + } + } else { + Iterator( + ( + value.asInstanceOf[Message], + (builder: Message.Builder, m: Message) => { + builder.setField(field, m) + () + })) + } + } + + /** + * `message` with every message inside it that `pf` matches replaced by what `pf` returns. A + * message nothing inside changed is returned as is. + */ + private def transformMessages(message: Message)( + pf: PartialFunction[Message, Message]): Message = + pf.applyOrElse( + message, + (m: Message) => { + var builder: Message.Builder = null + childMessages(m).foreach { case (child, set) => + val replaced = transformMessages(child)(pf) + if (replaced ne child) { + if (builder == null) { + builder = m.toBuilder + } + set(builder, replaced) + } + } + if (builder == null) m else builder.build() + }) + + /** The configs a Spark task would pass to native code with a plan. */ + private def nativeConf(): java.util.Map[String, String] = { + val conf = new java.util.HashMap[String, String]( + ConfigMap.parseFrom(CometExecIterator.serializeCometSQLConfs()).getEntriesMap) + // A Spark task passes the batch size apart from the other configs. + conf.put(CometConf.COMET_BATCH_SIZE.key, CometConf.COMET_BATCH_SIZE.get().toString) + conf + } + + private val singlePartition = PartitioningOuterClass.Partitioning + .newBuilder() + .setSinglePartition(PartitioningOuterClass.SinglePartition.newBuilder()) + .build() + + private class StageGraphBuilder(query: SparkPlan) { + val stages: ArrayBuffer[Stage] = ArrayBuffer.empty + private val stageIndexes = new IdentityHashMap[SparkPlan, Integer]() + + /** + * The index of the stage that runs the block of Comet operators rooted at `plan`, adding it + * first. + */ + def stageOf(plan: SparkPlan): Int = { + val known = stageIndexes.get(plan) + if (known != null) { + return known + } + val block = plan match { + case take: CometTakeOrderedAndProjectExec => return takeOrderedStage(take) + case native: CometNativeExec if native.serializedPlanOpt.isDefined => native + case other => fail(other, "it is not a Comet operator that runs in Rust") + } + val inputs = inputsOf(block).map { + case exchange: CometShuffleExchangeExec => + stageInput(exchange.child)(_.setPartitioning(partitioningOf(exchange))) + case broadcast => stageInput(broadcast.children.head)(_.setBroadcast(true)) + } + val blockPlan = Operator.parseFrom(block.serializedPlanOpt.plan.get) + + // A block that reads files runs one copy per partition, each naming its own files. + val (common, perPartition) = PlanDataInjector.findAllPlanData(block) + val partitionPlans = if (perPartition.isEmpty) { + Seq.empty + } else { + val base = new PlanDataInjector.CachedPlanData(blockPlan) + val counts = perPartition.map { case (key, data) => key -> data.length } + require( + counts.values.toSet.size == 1, + s"The scans of one block must have the same number of partitions, but found $counts") + (0 until counts.values.head).map { i => + val partitionData = perPartition.map { case (key, data) => key -> data(i) } + PlanDataInjector.injectPlanData(base, common, partitionData).toByteString + } + } + // The cluster reads it only for a stage that reads no exchange. A stage that does has as + // many partitions as its exchanges deliver. + val numPartitions = + if (partitionPlans.nonEmpty) partitionPlans.length + else block.outputPartitioning.numPartitions + // A stage that reads a range exchange sorts each of its partitions. A cluster that supports + // it merges them in this order instead of range-partitioning the exchange. Without it, or + // on a cluster that ignores it, the exchange is read as one partition. + val outputOrdering = if (inputs.exists(_.getPartitioning.hasRangePartition)) { + val orders = block.outputOrdering.map(QueryPlanSerde.exprToProto(_, block.output)) + if (orders.nonEmpty && orders.forall(_.isDefined)) orders.flatten else Nil + } else { + Nil + } + addStage(plan, blockPlan, inputs, partitionPlans, numPartitions, outputOrdering) + } + + /** An input of a stage that reads the stage running `producer`, as `configure` says. */ + private def stageInput(producer: SparkPlan)( + configure: StageInput.Builder => StageInput.Builder): StageInput = + configure(StageInput.newBuilder().setProducer(stageOf(producer))).build() + + /** + * A stage for `take` that reads its child's whole output as one partition and keeps the first + * rows in order, as Comet's final top-K plan does after its shuffle to one partition. A top-K + * of each partition first would send fewer rows through the exchange + * (https://github.com/apache/datafusion-comet/issues/7365). + */ + private def takeOrderedStage(take: CometTakeOrderedAndProjectExec): Int = { + val input = stageInput(take.child)(_.setPartitioning(singlePartition)) + // More than one input partition, so that the plan sorts across all of them. + val blockPlan = take + .finalNativePlan(inputPartitions = 2) + .getOrElse(fail(take, "its top-K plan cannot be built")) + addStage(take, blockPlan, Seq(input), Seq.empty, 1, Nil) + } + + private def addStage( + plan: SparkPlan, + blockPlan: Operator, + inputs: Seq[StageInput], + partitionPlans: Seq[ByteString], + numPartitions: Int, + outputOrdering: Seq[ExprOuterClass.Expr]): Int = { + findJvmCallback(blockPlan).foreach { name => + fail(plan, s"its native plan calls back into the JVM ($name)") + } + // Declared as Comet's shuffle writer declares its input, since the cluster casts the + // block's output to it as that writer does. + val outputSchema = schema2Proto( + plan.output.map(a => StructField(a.name, a.dataType, a.nullable, a.metadata))) + val stage = Stage + .newBuilder() + .setPlan(blockPlan) + .addAllInputs(inputs.asJava) + .addAllPartitionPlans(partitionPlans.asJava) + .setNumPartitions(numPartitions) + .addAllOutputSchema(outputSchema.asJava) + .addAllOutputOrdering(outputOrdering.asJava) + estimatedOutputBytes(plan).foreach(stage.setEstimatedOutputBytes) + stages += stage.build() + val index = stages.length - 1 + stageIndexes.put(plan, index) + index + } + + /** + * Spark's estimate of the size of `plan`'s output, from the logical plan of the nearest node + * in it that has one. A scheduler can use it to decide whether a join's build side is worth + * measuring before its probe side is shuffled. + */ + private def estimatedOutputBytes(plan: SparkPlan): Option[Long] = + plan + .collectFirst { case node if node.logicalLink.isDefined => node.logicalLink.get } + .flatMap(logical => Try(logical.stats.sizeInBytes).toOption) + .filter(_ >= 0) + .map(_.min(Long.MaxValue).toLong) + + /** + * The exchanges that feed `block`, in the order its native plan's input leaves read them: + * native shuffles and broadcasts, with a reused exchange resolved to the one it reuses. + */ + private def inputsOf(block: CometNativeExec): Seq[SparkPlan] = { + val inputs = ArrayBuffer.empty[SparkPlan] + def walk(plan: SparkPlan): Unit = plan match { + case reused: ReusedExchangeExec => + walk(reused.child) + case exchange: CometShuffleExchangeExec if exchange.shuffleType == CometNativeShuffle => + inputs += exchange + case broadcast: CometBroadcastExchangeExec => + inputs += broadcast + // These read or generate their rows inside the block. + case scan: CometNativeScanExec => + rejectDynamicPruning(scan, scan.partitionFilters) + case scan: CometIcebergNativeScanExec => + rejectDynamicPruning(scan, scan.runtimeFilters) + case _: CometRangeExec => + // These stand for rows that the JVM feeds to the plan, which a cluster has no way to do. + case _: CometScanWrapper | _: CometSinkPlaceHolder => + fail(plan, "its input comes from the JVM, which a cluster cannot read") + case _: CometNativeWriteExec | _: CometWriteFilesExec | _: CometIcebergWriteExec => + fail( + plan, + "Comet Native Executors do not support writes yet " + + "(https://github.com/apache/datafusion-comet/issues/7371)") + // Any other scan: its files are not described to the cluster. + case leaf: CometLeafExec => + fail(leaf, "Comet Native Executors do not support it yet") + // Part of the block, even when it has a serialized plan of its own: the conversion + // serializes some operators before their parent joins them into a larger block. + case native: CometNativeExec => + native.children.foreach(walk) + case other => fail(other, "Comet Native Executors do not support it yet") + } + walk(block) + inputs.toSeq + } + + /** + * A scan that prunes partitions with a dynamic filter lists its files only once the filter's + * subquery has run, which the driver cannot do before it builds the query. That includes the + * placeholder AQE leaves in a filter until it plans the subquery. + */ + private def rejectDynamicPruning(scan: SparkPlan, filters: Seq[Expression]): Unit = { + val dynamic = CometScanUtils + .filterUnusedDynamicPruningExpressions(filters) + .exists(_.exists { + case _: DynamicPruningExpression | _: DynamicPruning => true + case _ => false + }) + if (dynamic) { + fail( + scan, + "it prunes partitions with a dynamic filter (dynamic partition pruning), which " + + "Comet Native Executors do not support yet " + + "(https://github.com/apache/datafusion-comet/issues/7370)") + } + } + + private def partitioningOf( + exchange: CometShuffleExchangeExec): PartitioningOuterClass.Partitioning = { + def partitioning = PartitioningOuterClass.Partitioning.newBuilder() + exchange.outputPartitioning match { + case HashPartitioning(expressions, numPartitions) => + val keys = expressions.map { key => Review Comment: A shuffled join on `left("k") + 1 === right("k2")` passes planning. The cluster then rejects it at submission with "Hash partitioning keys must be columns". Could non-`Attribute` keys go through `fail` so the error names `CometExchange`? -- 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]
