Repository navigation
Conversation
…ture Part of apache#7288. The client submits a query to a cluster's scheduler over Arrow Flight and reads the result partitions from the executors that hold them. It is compiled only with the native-executors feature, which is off by default, so the standard build links no gRPC client. The wire format gains the result endpoints the client hands to the JVM and the stage fields the driver fills in.
Part of apache#7288. With spark.comet.nativeExecutors.enabled, CometRule replaces the plan Comet would execute with CometNativeExecutorsExec. The driver turns the plan into a stage graph, submits it to a cluster's scheduler through the JNI client, and Spark tasks read the result partitions from the executors over Arrow Flight. A query that cannot run entirely in the Comet pipeline fails, naming the operator. The two configs are internal, and the suite starts a cluster from COMET_NATIVE_EXECUTORS_BIN_DIR and skips without it.
…t in CI Part of apache#7288. Adds the contributor guide page for the driver, with how to run it and its current limitations. Registers the manual suite with the suite check, and lints and tests the native-executors feature in the Rust job.
viirya
left a comment
There was a problem hiding this comment.
Two points not tied to a line:
- The cluster-backed tests of
CometNativeExecutorsSuitewere run only against the cluster from #7289, never against the one in this repository. Since this PR's claim is that the driver works with that cluster, could you run the suite againstnative/native-executorsand report the result? The two clusters behave differently on range exchanges, on result serving and in their error messages, so a pass on #7289 doesn't carry over. - CI runs none of the new Scala code. The suite is in the manual list, and even its planning-only tests need a feature-built libcomet because
offloadcallsrequireNativeSupport()first.buildQueryandwithSubqueryResultsdon't need the native feature, though. A test that takes a normal Comet plan (AQE off) and checks the stage graph, the refusals, and the subquery-literal rewrite could run in the regular matrix.
| * 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 = |
There was a problem hiding this comment.
convert returns the plan unchanged when isCometLoaded(conf) is false or for a streaming plan. So with this config on, setting spark.comet.enabled=false or running any streaming query fails in buildQuery with "not a Comet operator that runs in Rust" instead of running in Spark. planOnlyApplies guards on isCometLoaded, isStreamingPlan and COMET_EXEC_ENABLED. Should this mirror it?
There was a problem hiding this comment.
Yes, it should. The offload now sits behind the same guards as the conversion: Comet loaded, exec enabled and not a streaming plan. CometRule shares them between plan-only mode and this. CometNativeExecutorsPlanSuite has a test that a query runs in Spark with the config on and Comet off.
| 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) |
There was a problem hiding this comment.
findAllPlanData is an execution-time hook. On each scan it calls ensureSubqueriesResolved() and sendDriverMetrics(), and it forces CometNativeScanExec.serializedPartitionData, which runs any scalar subquery in dataFilters and lists the files. Calling it from CometRule moves all of that into planning. For example, EXPLAIN select ... where v > (select avg(v) from t) now submits the subquery to the cluster, and it fails when the scheduler is down. The per-file partition plans also stay on the plan node for the Dataset's lifetime. Could planning only validate the graph, and the partition plans be built in doExecuteColumnar?
There was a problem hiding this comment.
Agreed. Planning now only validates the stage graph (CometNativeExecutors.validate), which lists no files and resolves no subquery. The query is built in doExecuteColumnar, where findAllPlanData runs. So EXPLAIN no longer submits the subquery or contacts the cluster, and the per-file plans no longer stay on the plan node. To do that the exec holds the Comet plan (nativePlan) instead of the serialized query.
| } | ||
| // 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 = |
There was a problem hiding this comment.
When a non-bucketed native scan lists no files (an empty table, or partition pruning that leaves nothing), findAllPlanData returns no data. This then falls back to UnknownPartitioning(0).numPartitions, and the cluster's build_stage rejects the stage with "The stage has no partitions". So select count(*) from t where p = 'missing' fails instead of returning 0. Sending a single partition with an empty file list would match Spark.
There was a problem hiding this comment.
Fixed for the native Parquet scan. A block whose scans list no files now gets one partition with an empty file list for each, so select count(*) from t where p = 'missing' returns 0. There is a test for it in CometNativeExecutorsPlanSuite and one on the cluster. I did not change the Iceberg scan.
| query, | ||
| subqueryResults.collect { case sub: ScalarSubquery => sub }) | ||
| val endpoints = | ||
| ResultEndpoints.parseFrom(new NativeExecutorsClient().submit(schedulerUrl, request)) |
There was a problem hiding this comment.
#7288 puts the wait inside a Spark task and cancels through CancelFlightInfo, so that "no driver thread blocks indefinitely". Here the driver thread stays in native code until the cluster finishes, and neither job cancellation nor an interrupt reaches it. Is moving the wait to the driver intentional? If so, the docs' cancellation item should also say that the driver thread stays blocked.
There was a problem hiding this comment.
Not a deliberate change from #7288. It is the first version's simplification: submit waits on the driver thread. The cancellation item in Current limitations now says that the driver thread stays blocked until the cluster finishes and that cancelling the job or interrupting the thread does not stop it. Moving the wait into a Spark task, together with CancelFlightInfo, belongs to #7372.
| * 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 = { |
There was a problem hiding this comment.
#7288 plans to drop DPP in this mode, since removing it never changes results. This refuses the query instead. DPP is on by default for any partitioned table joined to a filtered dimension, so refusing turns a common query shape into a failure. Was refusing chosen over dropping the filter?
There was a problem hiding this comment.
Refusing was the first version's scope. Dropping the filter means rebuilding the scan without it before its files are listed, and I kept that out of this PR. Queries with DPP now fail with a message that names the scan and links #7370, which is where dropping the filter, as #7288 plans, will land. Do you want it in this PR instead?
| "Cargo feature" | ||
|
|
||
| /** Whether the loaded libcomet has the client's native methods, which only the feature adds. */ | ||
| private lazy val hasNativeSupport: Boolean = { |
There was a problem hiding this comment.
NativeBase.isFeatureEnabled already does this for jemalloc and hdfs-opendal. Adding "native-executors" => cfg!(feature = "native-executors") to is_feature_enabled would avoid probing with closePartition(0) and matching on the UnsatisfiedLinkError message.
There was a problem hiding this comment.
Done. hasNativeSupport calls NativeBase.isFeatureEnabled("native-executors"), and is_feature_enabled reports the feature. The probe with closePartition(0) is gone.
| .endpoint | ||
| .into_iter() | ||
| .map(|endpoint| { | ||
| // A partition without a location is served by the scheduler itself. |
There was a problem hiding this comment.
The scheduler in this repository always sets a location and leaves DoGet unimplemented. So this fallback, and the NotFound/"expired" branch of read_error (line 68), only apply to the #7289 scheduler. I'd leave both out until that scheduler lands, together with the test cases built on its messages.
There was a problem hiding this comment.
Removed. A partition without a location is now an error, and so is the NotFound branch of read_error with its test case. I kept the Failed to open partition file ... NotFound case, since the executor in this repository returns it once its retention has passed.
| // Spark would sample for them. A cluster that can read the producer's partitions as they are, | ||
| // sort each, and merge them in this order may do that; one that ignores this field reads the | ||
| // exchange as one partition, which the block then sorts. | ||
| repeated spark.spark_expression.Expr output_ordering = 6; |
There was a problem hiding this comment.
Nothing in this repository reads output_ordering or estimated_output_bytes. Could these fields, and the driver code that fills them (outputOrdering in stageOf and estimatedOutputBytes), come with the cluster change that uses them? Until then, these comments describe behavior the cluster here doesn't have.
There was a problem hiding this comment.
Removed both fields, the driver code that fills them, and the mentions in the docs. They can come with the cluster change that reads them.
| cargo clippy --color=never -p datafusion-comet --all-targets --features python-udf -- -D warnings | ||
|
|
||
| # The JNI client for Comet Native Executors is off by default, so nothing above compiles it. | ||
| - name: Check native executors feature |
There was a problem hiding this comment.
Nit: this adds another full build of datafusion-comet. Folding it into the step above as --features python-udf,native-executors would lint both features in one build.
There was a problem hiding this comment.
Folded into the python-udf step, which now lints --features python-udf,native-executors in one build.
rich7420
left a comment
There was a problem hiding this comment.
@andygrove thanks for the patch!
| let info = client | ||
| .get_flight_info(FlightDescriptor::new_cmd(query)) | ||
| .await | ||
| .map_err(|e| client_error("the query failed", e))?; |
There was a problem hiding this 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?
There was a problem hiding this comment.
Added to Current limitations. A failure on the cluster reaches Spark as a CometNativeException with the cluster's message, and the page uses the ANSI division by zero as the example. Rebuilding Spark's exception needs the cluster to send a structured error instead of Ballista's Debug string, which I would do as a follow-up with the rest of the error handling rather than here.
| * 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) |
There was a problem hiding this 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?
There was a problem hiding this comment.
Good catch. checkOnCluster now takes Spark's answer with COMET_ENABLED=false, as the TPC-H test does.
| def partitioning = PartitioningOuterClass.Partitioning.newBuilder() | ||
| exchange.outputPartitioning match { | ||
| case HashPartitioning(expressions, numPartitions) => | ||
| val keys = expressions.map { key => |
There was a problem hiding this 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?
There was a problem hiding this comment.
Fixed. A hash exchange whose key is not a column is refused at planning, naming the exchange, with the reason that a cluster hashes columns. There is a test in CometNativeExecutorsPlanSuite.
…im what no cluster here uses Review feedback on the driver: - Planning only validates the stage graph. The files a scan reads are listed, and the subqueries in its filters run, when the query executes, so EXPLAIN no longer contacts the cluster and the per-file plans no longer stay on the plan node. - A scan that lists no files gets one empty partition, so a query over an empty table or a fully pruned partition returns no rows instead of failing in the scheduler. - The offload follows the conversion's own guards: Comet loaded, exec enabled, not streaming. - A hash exchange keyed by an expression is refused naming the exchange, since the cluster hashes columns. - Remove the Stage fields output_ordering and estimated_output_bytes, which nothing in this repository reads, and the client's handling of a scheduler that serves partitions itself. - Ask libcomet whether it was built with the feature through NativeBase.isFeatureEnabled. - Add CometNativeExecutorsPlanSuite, which checks the stage graph, the refusals and the subquery rewrite in the regular CI matrix, and compare the cluster suite with Spark's answer with Comet off. - Document the blocked driver thread on cancellation and the generic exception for cluster errors, and say which TPC-H tests need the cluster of apache#7289.
|
On the two points that were not tied to a line: I ran I added |
rich7420
left a comment
There was a problem hiding this comment.
@andygrove thanks for the update!!!
| * @throws CometRuntimeException | ||
| * naming the first operator that cannot run on a cluster | ||
| */ | ||
| def validate(plan: SparkPlan): Unit = stageGraph(plan, withFiles = false) |
There was a problem hiding this 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.
| override protected def doExecuteColumnar(): RDD[ColumnarBatch] = { | ||
| NativeExecutorsClient.requireNativeSupport() | ||
| val request = CometNativeExecutors.withSubqueryResults( | ||
| CometNativeExecutors.buildQuery(nativePlan).toByteArray, |
There was a problem hiding this 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?
| 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 |
There was a problem hiding this 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.
| 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( |
comphead
left a comment
There was a problem hiding this comment.
I checked the fixes from the earlier rounds at f1819d5ba5 and they are in. With the config off, CometRule behaves as before. The comments below are about task-side cancellation, result expiry, what EXPLAIN shows, and tests that could run without a cluster.
| context.addTaskCompletionListener[Unit](_ => batches.close()) | ||
| batches |
There was a problem hiding this 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.
| 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 }) |
There was a problem hiding this 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.
| 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") { |
There was a problem hiding this 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.
| override def stringArgs: Iterator[Any] = | ||
| Iterator(output, schedulerUrl) |
There was a problem hiding this 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.
| 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) } |
There was a problem hiding this 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.
| - 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 |
There was a problem hiding this 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.
Which issue does this PR close?
Part of #7288.
Rationale for this change
Comet Native Executors run whole Spark queries on a cluster of
comet-schedulerandcomet-executorprocesses with no JVM in the data plane. The cluster is innative/native-executors. This PR adds the Spark driver side: the driver plans a query with Catalyst and Comet, turns the plan into a graph of stages, and submits it to the scheduler's Arrow Flight service. Spark tasks then read the result partitions from the executors over Flight.The driver never links Ballista. Its only new native dependencies are
arrow-flightandtonic, and they are optional. The feature is off by default at every level: thenative-executorsCargo feature is off,spark.comet.nativeExecutors.enabledis false, and both configs are internal. With the config off,CometRulebehaves as before.This PR is split out of #7289. The cluster-side features some queries need, such as sorting the partitions of a range exchange in parallel and moving a join whose build side turns out small, are in that PR and not here. This PR is not expected to run the TPC-H queries yet: the cluster in this repository fails most of the ones that join. Nothing here depends on Ballista's
comet-nativebranch.What changes are included in this PR?
One commit for each part, and a commit for the review feedback.
native_executors.protogainsResultEndpointsandResultEndpoint.native/coregains thenative-executorsCargo feature with optionalarrow-flightandtonicdependencies,native_executors/client.rs(submit a query, open and read a result partition, close it), andexport_batchbecomespub(crate)for the client.is_feature_enabledreports the feature.native/Cargo.lockis regenerated.CometNativeExecutorsExec.scalahas the planner. Planning validates the stage graph and refuses a query that cannot run entirely in the Comet pipeline, naming the operator. The query itself is built when it executes: that lists the files each scan reads, runs the subqueries in the scans' filters, and writes scalar subquery results into the plan as literals. A scan that lists no files gets one empty partition. The exec then reads the result partitions.NativeExecutorsClient.scalahas the JNI declarations and a clear error when libcomet lacks the feature.CometRulecalls the planner when the config is on and Comet converts the plan.CometConfgetsspark.comet.nativeExecutors.enabledandspark.comet.nativeExecutors.scheduler.url.CometNativeExecutorsPlanSuitechecks the stage graph, the refusals and the subquery rewrite without a cluster, in the regular CI matrix.CometNativeExecutorsSuiteruns queries on a cluster and is manual.pom.xmlexempts the JNI declarations ofNativeExecutorsClientfrom the unused-parameter check of thestrict-warningsprofile, as it does forNative.native_executors.md, with how to run the driver and its current limitations, and the one sentence ofnative/native-executors/README.mdthat said the driver was in progress. The CI parts are a clippy step and a nextest step for thenative-executorsfeature in the Rust test action, and the manual suite's entry indev/ci/check-suites.py.How are these changes tested?
Rust:
cargo clippy -p datafusion-comet --all-targets --features python-udf,native-executors -- -D warnings,cargo fmt --check, andcargo test -p datafusion-comet --features native-executors --lib native_executors feature_enabled.Scala, with JDK 17 on the default profile:
./mvnw spotless:check,./mvnw test-compile -DskipTests,CometNativeExecutorsPlanSuite,CometExecRuleSuiteandCometIcebergWriteDetectionSuite. Markdown passesprettier --check, anddev/ci/check-suites.pypasses.The cluster-backed tests of
CometNativeExecutorsSuiteare manual, since they need the cluster binaries inCOMET_NATIVE_EXECUTORS_BIN_DIR. They ran against the cluster built fromnative/native-executorsin this PR, with a libcomet built from this PR with thenative-executorsfeature, on TPC-H data at scale factor 1. All 15 tests that are not TPC-H queries pass, and so do TPC-H q1, q6, q11 and q15. The other 18 TPC-H queries join, and fail while the scheduler plans them (CometFragmentExec input N has K partitions, expected 1). They need the cluster of #7289.