comphead commented on code in PR #7544: URL: https://github.com/apache/datafusion-comet/pull/7544#discussion_r4532817216
########## 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 Review Comment: This iterator is where the task's data comes from, and as far as I can tell nothing above it checks whether the task was killed. So when Spark kills the task (a cancelled job, or a speculative copy that lost), the task keeps reading its partition to the end. Spark's own source RDDs wrap their iterator in `InterruptibleIterator`, and `CometBlockStoreShuffleReader` does the same for Comet's shuffle reads. Could this return `new InterruptibleIterator(context, batches)`? A read that is blocked inside `nextBatch` still waits for its batch, but the task would stop at the next batch boundary. The Cancellation item in the docs could then say that too. ########## native/core/src/execution/native_executors/client.rs: ########## @@ -0,0 +1,263 @@ +// 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: 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) if status.code() == Code::Internal => { + let message = status.message(); + message.contains("Failed to open partition file") && message.contains("NotFound") + } + _ => 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))?; + let endpoints = info + .endpoint + .into_iter() + .map(|endpoint| { + let location = endpoint + .location + .first() + .map(|location| location.uri.clone()) + .ok_or_else(|| client_error("bad result", "a partition has no location"))?; + let ticket = endpoint + .ticket + .map(|ticket| ticket.ticket.to_vec()) + .ok_or_else(|| client_error("bad result", "a partition has no ticket"))?; + Ok(ResultEndpoint { location, ticket }) Review Comment: The scheduler sets `expiration_time` on each endpoint (`results_expire_at` in `frontend.rs`), and the cluster README tells clients to rely on it. This conversion drops it. That leaves `read_error` to spot an expired partition by finding `Failed to open partition file` and `NotFound` in the text of Ballista's error. That text can change with the Ballista pin, and the unit test hardcodes it, so it would not catch the change. Could `ResultEndpoint` carry the expiration time? Then a read that fails after it could give the retention hint without parsing the message, and a task that starts after it could fail with the hint before it connects. ########## spark/src/test/scala/org/apache/comet/exec/CometNativeExecutorsSuite.scala: ########## @@ -0,0 +1,482 @@ +/* + * 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 + * }}} + * + * The TPC-H tests need the cluster of https://github.com/apache/datafusion-comet/pull/7289. On + * the cluster in this repository, q1, q6, q11 and q15 pass and the other queries, which join, + * fail while the scheduler plans them. + * + * `CometNativeExecutorsPlanSuite` checks the stage graphs without a cluster, and runs in CI. + */ +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, sparkAnswer(query), ordered) + + /** Spark's own answer, with Comet off, so that a bug the cluster shares with Comet shows. */ + private def sparkAnswer(query: => DataFrame): Seq[Row] = { + // withSQLConf returns Unit on Spark 3.x, so the rows come out through a var. + var rows: Seq[Row] = Nil + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + rows = query.collect().toSeq + } + rows + } + + /** Checks that the cluster runs `query` and returns `expected`, in order if `ordered`. */ + private def checkClusterReturns( + query: => DataFrame, + expected: Seq[Row], + ordered: Boolean): Unit = onCluster { + // A new Dataset, since a Dataset keeps the plan it was first run with. + val df = query + if (ordered) { + assert(df.collect().toSeq == expected, "the cluster returned other rows or another order") + } else { + checkAnswer(df, expected) + } + val plan = df.queryExecution.executedPlan + assert(collect(plan) { case exec: CometNativeExecutorsExec => exec }.size == 1, plan) + } + + /** A Parquet table with several files, read as several partitions. */ + private def withTable(body: String => Unit): Unit = withTempPath { dir => + spark + .range(0, 10000, 1, 8) + .selectExpr("id", "id % 7 AS k", "CAST(id AS DOUBLE) * 1.5 AS v") + .write + .parquet(dir.getCanonicalPath) + withSQLConf(SQLConf.FILES_MAX_PARTITION_BYTES.key -> "16384") { + body(dir.getCanonicalPath) + } + } + + nativeExecutorsTest("scan, filter and project") { + withTable { path => + checkOnCluster(spark.read.parquet(path).where("k > 2").selectExpr("id", "v * 2 AS w")) + } + } + + nativeExecutorsTest("aggregate across a hash shuffle") { + withTable { path => + checkOnCluster(spark.read.parquet(path).groupBy("k").agg(sum("v"), count("*"))) + } + } + + nativeExecutorsTest("aggregate across a hash shuffle without AQE") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withTable { path => + checkOnCluster(spark.read.parquet(path).groupBy("k").agg(sum("v"), max("id"))) + } + } + } + + nativeExecutorsTest("global aggregate across a single-partition shuffle") { + withTable { path => + checkOnCluster(spark.read.parquet(path).agg(sum("id"), max("v"), count("*"))) + } + } + + Seq(true, false).foreach { sortMerge => + val join = if (sortMerge) "sort-merge" else "hash" + nativeExecutorsTest(s"shuffled $join join") { + withSQLConf( + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.PREFER_SORTMERGEJOIN.key -> sortMerge.toString) { + withTable { path => + val left = spark.read.parquet(path).select("id", "k") + val right = spark.read + .parquet(path) + .where("id < 10") + .select(col("k").as("k2"), col("v")) + checkOnCluster(left.join(right, left("k") === right("k2")).select("id", "v")) + } + } + } + } + + /** + * TPC-H's `lineitem`: the table under `SPARK_TPCH_DATA` when it is set, as for + * `CometTPCHQuerySuite`, and otherwise generated data with the same schema. + */ + private def withLineitem(body: => Unit): Unit = sys.env.get("SPARK_TPCH_DATA") match { + case Some(path) => + spark.read.parquet(s"$path/lineitem").createOrReplaceTempView("lineitem") + withSQLConf(SQLConf.FILES_MAX_PARTITION_BYTES.key -> "32m")(body) + case None => + withTempPath { dir => + spark + .range(0, 100000, 1, 8) + .selectExpr( + "CAST(id % 50 + 1 AS DECIMAL(15, 2)) AS l_quantity", + "CAST((id % 1000) * 1.37 + 900 AS DECIMAL(15, 2)) AS l_extendedprice", + "CAST((id % 11) / 100.0 AS DECIMAL(15, 2)) AS l_discount", + "CAST((id % 9) / 100.0 AS DECIMAL(15, 2)) AS l_tax", + "ELT(CAST(id % 3 AS INT) + 1, 'A', 'N', 'R') AS l_returnflag", + "ELT(CAST(id % 2 AS INT) + 1, 'F', 'O') AS l_linestatus", + "DATE_ADD(DATE '1992-01-01', CAST(id % 2500 AS INT)) AS l_shipdate") + .write + .parquet(dir.getCanonicalPath) + spark.read.parquet(dir.getCanonicalPath).createOrReplaceTempView("lineitem") + withSQLConf(SQLConf.FILES_MAX_PARTITION_BYTES.key -> "131072")(body) + } + } + + private val q1Aggregate = """ + |select + | l_returnflag, + | l_linestatus, + | sum(l_quantity) as sum_qty, + | sum(l_extendedprice) as sum_base_price, + | sum(l_extendedprice * (1 - l_discount)) as sum_disc_price, + | sum(l_extendedprice * (1 - l_discount) * (1 + l_tax)) as sum_charge, + | avg(l_quantity) as avg_qty, + | avg(l_extendedprice) as avg_price, + | avg(l_discount) as avg_disc, + | count(*) as count_order + |from + | lineitem + |where + | l_shipdate <= date '1998-12-01' - interval '90' day + |group by + | l_returnflag, + | l_linestatus""".stripMargin + + nativeExecutorsTest("TPC-H q1 without its order by") { + withLineitem { + checkOnCluster(sql(q1Aggregate)) + } + } + + nativeExecutorsTest("TPC-H q1") { + withLineitem { + checkOnCluster( + sql(s"$q1Aggregate\norder by\n l_returnflag,\n l_linestatus"), + ordered = true) + } + } + + private val tpchTables = + Seq("customer", "lineitem", "nation", "orders", "part", "partsupp", "region", "supplier") + + /** + * The 22 TPC-H queries, on the tables under `SPARK_TPCH_DATA`. `TPCH_QUERIES`, a + * comma-separated list of query numbers, runs only those. + */ + private val tpchQueries = { + val only = sys.env.get("TPCH_QUERIES").map(_.split(",").map(_.trim.toInt).toSet) + (1 to 22).filter(q => only.forall(_.contains(q))) + } + + tpchQueries.foreach { q => + nativeExecutorsTest(s"TPC-H query $q") { + val data = sys.env.get("SPARK_TPCH_DATA") + assume(data.isDefined, "SPARK_TPCH_DATA is not set") + tpchTables.foreach { table => + spark.read.parquet(s"${data.get}/$table").createOrReplaceTempView(table) + } + val file = + new File(getClass.getResource("/").toURI.resolve("../../../benchmarks/tpc/queries/tpch/")) + val text = new String(Files.readAllBytes(new File(file, s"q$q.sql").toPath), "UTF-8") + val statements = text + .split("\n") + .filterNot(_.trim.startsWith("--")) + .mkString("\n") + .split(";") + .map(_.trim) + .filter(_.nonEmpty) + // q15 creates a view, queries it and drops it. + val (setup, rest) = statements.span(!_.toLowerCase.startsWith("select")) + val query = rest.head + val teardown = rest.tail + // The tables are temporary views, which only a temporary view can read. + setup.map(_.replaceFirst("(?i)^create view", "create temp view")).foreach(sql) + try { + // Spark's own answer. Comet on the driver would run the whole query in one JVM, in the + // test's small native memory pool, and Spark's broadcast relations don't fit in the test + // JVM's heap either. + // withSQLConf returns Unit on Spark 3.x, so the rows come out through a var. + var expected: Seq[Row] = Nil + withSQLConf( + CometConf.COMET_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + expected = sql(query).collect().toSeq + } + checkClusterReturns(sql(query), expected, query.toLowerCase.contains("order by")) + } finally { + teardown.foreach(sql) + } + } + } + + nativeExecutorsTest("broadcast hash join") { + withTable { path => + val left = spark.read.parquet(path).select("id", "k") + val right = spark.read.parquet(path).where("id < 10").select(col("k").as("k2"), col("v")) + checkOnCluster(left.join(right, left("k") === right("k2")).select("id", "v")) + } + } + + nativeExecutorsTest("a scan that lists no files returns no rows") { + withTempPath { dir => + spark + .range(0, 100) + .selectExpr("id", "id % 7 AS k") + .write + .partitionBy("k") + .parquet(dir.getCanonicalPath) + // No partition has k = 99, so the scan has no files to read. + def table = spark.read.parquet(dir.getCanonicalPath).where("k = 99") + checkOnCluster(table) + checkOnCluster(table.groupBy().count()) + } + } + + nativeExecutorsTest("scalar subquery") { + withTable { path => + spark.read.parquet(path).createOrReplaceTempView("t") + checkOnCluster(sql("select id, v from t where v > (select avg(v) from t where k = 3)")) + } + } + + nativeExecutorsTest("a command's plan stays in Spark") { + withTable { path => + onCluster { + withTempView("t") { + spark.read.parquet(path).createOrReplaceTempView("t") + // Explaining a command explains its result, which isn't a query to run on the cluster. + val df = sql("create or replace temp view v as select * from t") + assert( + df.queryExecution.explainString(org.apache.spark.sql.execution.SimpleMode).nonEmpty) + } + } + } + } + + /** The messages of the exception that planning `df` on the cluster throws and its causes. */ + private def refusalMessages(df: => DataFrame): Seq[String] = { + // withSQLConf returns Unit on Spark 3.x, so the messages come out through a var. + var messages: Seq[String] = Nil + onCluster { + // A new Dataset, since a Dataset keeps the plan it was first run with. + val error = intercept[Exception](df.queryExecution.executedPlan) + messages = Iterator + .iterate[Throwable](error)(_.getCause) + .takeWhile(_ != null) + .map(e => String.valueOf(e.getMessage)) + .toSeq + } + messages + } + + /** Checks that the cluster refuses `df`, naming `operator` and giving `reason`. */ + private def checkRefused(df: => DataFrame, operator: String, reason: String): Unit = { + val messages = refusalMessages(df) + assert( + messages.exists(message => + message.contains("Comet Native Executors cannot run this query") && + message.contains(operator) && message.contains(reason)), + messages) + } + + nativeExecutorsTest("a query the cluster cannot run fails, naming the operator") { + withTable { path => + // With Comet's sort off, the sort stays in Spark. + withSQLConf(CometConf.COMET_EXEC_SORT_ENABLED.key -> "false") { + checkRefused(spark.read.parquet(path).orderBy("v"), "Sort", "") + } + } + } + + // The tests below fail while planning, so they need no cluster. + + test("a scan that stays on the JVM fails, naming the operator that reads from it") { Review Comment: This test checks the same refusal as `a sort that stays in Spark is refused, naming the operator` in `CometNativeExecutorsPlanSuite`, and with `reason` empty it checks less. The two tests below need no cluster either. They need only a libcomet built with the feature, because they go through `CometRule`, and `offload` calls `requireNativeSupport()` first. Through `CometNativeExecutors.validate`, as the plan suite's `refusal` does, the refusals of the JVM scan and of dynamic pruning would run in the regular CI matrix. Could they move there, and this test be dropped? `withTable` is also the same in both suites. ########## 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) Review Comment: `nativePlan` is not a child, so `EXPLAIN` shows a bare `CometNativeExecutors` leaf and nothing of the stages the cluster runs. `InMemoryTableScanExec` is built the same way, and it overrides `innerChildren` as `Seq(relation) ++ super.innerChildren` so that explain prints the cached plan under the leaf. Would `override def innerChildren: Seq[QueryPlan[_]] = Seq(nativePlan) ++ super.innerChildren` work here? I have not checked how formatted explain numbers the operators inside it. ########## .github/actions/rust-test/action.yaml: ########## @@ -106,6 +108,13 @@ runs: export LD_LIBRARY_PATH=${JAVA_HOME}/lib/server:${LD_LIBRARY_PATH} RUST_BACKTRACE=1 cargo nextest run -p datafusion-comet --lib --features python-udf python_udf + - name: Run native executors client tests + shell: bash + run: | + cd native + export LD_LIBRARY_PATH=${JAVA_HOME}/lib/server:${LD_LIBRARY_PATH} + RUST_BACKTRACE=1 cargo nextest run -p datafusion-comet --lib --features native-executors native_executors Review Comment: Like the clippy step before it was folded, this builds the `datafusion-comet` test binary once more for one feature. `cargo nextest run -p datafusion-comet --lib --features python-udf,native-executors python_udf native_executors` in the step above would run both sets of tests from one build. I have not run it. ########## native/core/src/execution/native_executors/client.rs: ########## @@ -0,0 +1,263 @@ +// 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: 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) if status.code() == Code::Internal => { + let message = status.message(); + message.contains("Failed to open partition file") && message.contains("NotFound") + } + _ => 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))?; + let endpoints = info + .endpoint + .into_iter() + .map(|endpoint| { + let location = endpoint + .location + .first() + .map(|location| location.uri.clone()) + .ok_or_else(|| client_error("bad result", "a partition has no location"))?; + let ticket = endpoint + .ticket + .map(|ticket| ticket.ticket.to_vec()) + .ok_or_else(|| client_error("bad result", "a partition has no ticket"))?; + Ok(ResultEndpoint { location, ticket }) + }) + .collect::<CometResult<Vec<_>>>()?; + Ok(ResultEndpoints { endpoints }) +} + +/// Runs a query: `query` is a serialized `QueryRequest`, which this fills in with the driver's +/// Comet version and plan format. Returns a serialized `ResultEndpoints`. +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_NativeExecutorsClient_submit( + e: EnvUnowned, + _class: JClass, + scheduler_url: JString, + query: JByteArray, +) -> jbyteArray { + try_unwrap_or_throw(&e, |env| { + let scheduler_url = scheduler_url + .try_to_string(env) + .map_err(|e| client_error("bad scheduler URL", e))?; + let query = env.convert_byte_array(query)?; + let endpoints = get_runtime().block_on(submit(&scheduler_url, query))?; + Ok(env + .byte_array_from_slice(&endpoints.encode_to_vec())? + .into_raw()) + }) +} + +/// Starts reading a result partition from the executor at `location`. Returns a handle for +/// `nextBatch` and `closePartition`. +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_NativeExecutorsClient_openPartition( + e: EnvUnowned, + _class: JClass, + location: JString, + ticket: JByteArray, +) -> jlong { + try_unwrap_or_throw(&e, |env| { + let location = location + .try_to_string(env) + .map_err(|e| client_error("bad location", e))?; + let ticket = env.convert_byte_array(ticket)?; + let stream = get_runtime().block_on(async { + connect(&location) + .await? + .do_get(Ticket::new(ticket)) + .await + .map_err(|e| read_error(&format!("cannot read a partition from {location}"), e)) + })?; + Ok(Box::into_raw(Box::new(stream)) as jlong) + }) +} + +/// Moves the partition's next batch into the Arrow C Data Interface structs at `array_addrs` +/// and `schema_addrs`, one pair per column. Returns its row count, or -1 at the end. +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_NativeExecutorsClient_nextBatch( + e: EnvUnowned, + _class: JClass, + handle: jlong, + array_addrs: JLongArray, + schema_addrs: JLongArray, +) -> jlong { + try_unwrap_or_throw(&e, |env| { + // SAFETY: `handle` comes from `openPartition` and hasn't been closed. + let stream = unsafe { &mut *(handle as *mut FlightRecordBatchStream) }; + let Some(batch) = get_runtime().block_on(stream.next()) else { + return Ok(-1); + }; + let batch = batch.map_err(|e| read_error("cannot read a partition", e))?; + let array_addrs = unsafe { array_addrs.get_elements(env, ReleaseMode::NoCopyBack)? }; + let schema_addrs = unsafe { schema_addrs.get_elements(env, ReleaseMode::NoCopyBack)? }; + // SAFETY: the JVM allocated one array and one schema struct per column. + unsafe { export_batch(&array_addrs, &schema_addrs, batch, false) } Review Comment: These lines are the body of `prepare_output` in `jni_api.rs`, which `Native.takeShuffleBatch` already calls as `prepare_output(env, array_addrs, schema_addrs, batch, false)`. If `prepare_output` became `pub(crate)` and this called it, `export_batch` could stay private and the two paths would share one helper. -- 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]
