You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/03 11:48:05 UTC
(datafusion-comet) 01/01: feat: add a native Comet version of RangeExec (#6553)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6553-93d1189beb2a2da1c4e2663de131c61713b51e48
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 1dcea020b0c7a01dd39bb0373d701767ad15a13b
Author: Andy Grove <ag...@apache.org>
AuthorDate: Sat Oct 3 11:47:42 2026 +0000
feat: add a native Comet version of RangeExec (#6553)
* feat: add an experimental Comet version of RangeExec
Add CometRangeExec, which writes the values of spark.range and SQL range()
straight into Arrow batches on the JVM and feeds them to the native plan, so
the operators above a range can run natively without converting Spark rows.
Partition bounds and values follow Spark's generated code for RangeExec,
including its behavior where the arithmetic overflows.
Off by default behind spark.comet.exec.range.enabled. When the Spark-to-Arrow
conversion is also enabled, CometRangeExec takes precedence for Range.
Adds CometRangeExecSuite and CometRangeBenchmark.
* docs: list the experimental Comet RangeExec on the operators page
* feat: generate Range values in native code
Add a native RangeExec to the operators crate, planned from a new RangeScan
protobuf message. Each task computes its partition's bounds from the
partition index in i128 arithmetic and walks them with the loop from Spark's
generated code for RangeExec, so the values match Spark's default path,
overflow cases included. On the JVM, CometNativeRangeExec is a CometLeafExec
that reports Spark's partitioning, so the native plan runs one task per slice
with no JVM input.
The internal spark.comet.exec.range.native.enabled (default true) switches
between this and the JVM generator. CometRangeExecSuite runs every test
against both, and CometRangeBenchmark has an arm for each.
* refactor: keep only the native Range generator
Drop the JVM generator (RangeArrowReader) and the internal switch between the
two, and name the native operator CometRangeExec. The native generator was
faster on every benchmark case.
* fix: keep Range on Spark where its interpreted path disagrees with codegen
With whole-stage codegen disabled, Spark runs the interpreted
RangeExec.doExecute, which returns different rows from its generated code
when that code's arithmetic overflows. Comet follows the generated code, so
such ranges now fall back when codegen is off.
Also list CometRange in the operator reference, mark originalPlan transient
as the other native leaves do, and test a global aggregate over an empty
range.
* fix: fail a Range with no slices as Spark does
Spark does not validate a range's slice count. A non-empty range with fewer
than one slice fails when Spark runs it, while Comet would plan no tasks and
return no rows, so such ranges now stay on Spark.
The Spark-to-Arrow conversion now steps aside for Range only where
CometRangeExec supports the range, so a range it declines can still use that
conversion.
* refactor: simplify CometRangeExec and generalize leaf precedence
- Build the native range on DataFusion's LazyMemoryExec with a
LazyBatchGenerator, as DataFusion's own generate_series does, instead of a
hand-written ExecutionPlan and stream. This also gives it cooperative
yielding. The loop now reads like Spark's doProduce, and each batch is sized
to the values left.
- In CometExecRule, a leaf with an enabled Comet operator now tries that
operator before the Spark-to-Arrow conversion, which becomes the fallback
for a leaf the operator declines. This replaces the Range special case in
shouldApplySparkToColumnar and records the decline reason.
- Share the codegen-disabled check, including the NO_CODEGEN factory mode on
Spark 3.5 and later, between CometRangeExec and the hash aggregate.
- Restore Spark's RangeExec where no native operator consumes the range,
since generating it natively would only add a conversion back to rows.
- Drop overrides CometExec already provides, and duplicate tests and
benchmark validation runs.
* refactor: keep overflowing ranges on Spark and generate the rest directly
Spark's generated code for RangeExec and its interpreted path, which runs
when whole-stage codegen is off or the generated code is too large, disagree
only for a range whose arithmetic overflows. Such ranges now stay on Spark
whatever the codegen setting, so Comet no longer has to predict which path
Spark takes. The check is exact: an element count that does not fit in a long,
or a batch span that does not.
Every range Comet runs now has the values start + i * step in both of
Spark's paths, so the native generator drops the emulation of Spark's
1000-value batches and computes them directly.
The codegen-disabled helper shared with the hash aggregate is no longer
needed, so the aggregate is back to main's version. The config and docs no
longer call the feature experimental. It stays off by default because it can
be slower than Spark under cheap operators.
---
.github/workflows/pr_build_linux.yml | 1 +
.github/workflows/pr_build_macos.yml | 1 +
docs/source/user-guide/latest/operators.md | 16 +-
.../user-guide/latest/understanding-comet-plans.md | 19 +-
native/Cargo.lock | 1 +
native/core/src/execution/jni_api.rs | 1 +
native/core/src/execution/planner.rs | 19 +-
.../src/execution/planner/operator_registry.rs | 1 +
native/operators/Cargo.toml | 1 +
native/operators/src/lib.rs | 2 +
native/operators/src/range.rs | 260 +++++++++++++++++++++
native/proto/src/proto/operator.proto | 12 +
.../main/scala/org/apache/comet/CometConf.scala | 9 +
.../org/apache/comet/rules/CometExecRule.scala | 17 ++
.../rules/EliminateRedundantTransitions.scala | 10 +-
.../apache/spark/sql/comet/CometRangeExec.scala | 127 ++++++++++
.../apache/comet/exec/CometRangeExecSuite.scala | 233 ++++++++++++++++++
.../spark/sql/benchmark/CometRangeBenchmark.scala | 133 +++++++++++
18 files changed, 844 insertions(+), 19 deletions(-)
diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml
index d85eb14d2c..57c695a271 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.github/workflows/pr_build_linux.yml
@@ -512,6 +512,7 @@ jobs:
org.apache.comet.exec.CometExecSuite
org.apache.comet.exec.CometTopKSuite
org.apache.comet.exec.CometEmptyRelationExecSuite
+ org.apache.comet.exec.CometRangeExecSuite
org.apache.comet.exec.CometInMemoryCacheSuite
org.apache.comet.exec.CometInMemoryCachePruningSuite
org.apache.comet.exec.CometInMemoryCacheKryoSuite
diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml
index 02d3077e4f..51b6d493cb 100644
--- a/.github/workflows/pr_build_macos.yml
+++ b/.github/workflows/pr_build_macos.yml
@@ -216,6 +216,7 @@ jobs:
org.apache.comet.exec.CometExecSuite
org.apache.comet.exec.CometTopKSuite
org.apache.comet.exec.CometEmptyRelationExecSuite
+ org.apache.comet.exec.CometRangeExecSuite
org.apache.comet.exec.CometInMemoryCacheSuite
org.apache.comet.exec.CometInMemoryCachePruningSuite
org.apache.comet.exec.CometInMemoryCacheKryoSuite
diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md
index 88fdc30ddc..676f295ccf 100644
--- a/docs/source/user-guide/latest/operators.md
+++ b/docs/source/user-guide/latest/operators.md
@@ -53,19 +53,19 @@ omitted from the tables below and may be reconsidered based on demand:
- **Structured Streaming operators** (`StateStoreSaveExec`, `StateStoreRestoreExec`, `StreamingSymmetricHashJoinExec`, and similar): Comet targets batch execution.
- **Cartesian / cross joins** (`CartesianProductExec`): rare and expensive, with little acceleration benefit.
-- **Range generation** (`RangeExec`): niche leaf operator.
- **Pickled (non-Arrow) Python UDFs** (`BatchEvalPythonExec`): Comet accelerates Arrow-based Python UDFs only ([#4234](https://github.com/apache/datafusion-comet/pull/4234)).
- **Typed Dataset operators** (`DeserializeToObjectExec`, `SerializeFromObjectExec`, `MapElementsExec`, `MapPartitionsExec`, `MapGroupsExec`, `CoGroupExec`, `AppendColumnsExec`, and similar): produced by `map`, `mapPartitions`, `groupByKey`, `cogroup`, and other typed `Dataset` transformations. They exist to run user JVM functions on JVM objects, which Comet cannot do natively.
## Scans
-| Operator | Status | Notes |
-| ----------------------- | ------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| `FileSourceScanExec` | ✅ | Parquet only. Some types and configurations fall back. See [Parquet Scan Compatibility](compatibility/scans.md). |
-| `BatchScanExec` | ✅ | Apache Iceberg Parquet scans run natively. Native CSV scans are experimental and disabled by default. DataSource V2 Parquet scans are not accelerated. See [Parquet Scan Compatibility](compatibility/scans.md) and the [Iceberg Guide](iceberg.md). |
-| `LocalTableScanExec` | ⚠️ | Disabled by default; there is no acceleration advantage and this operator is typically only used in test code. Can be opted into via config ([#4393](https://github.com/apache/datafusion-comet/pull/4393)). |
-| `EmptyRelationExec` | ✅ | Spark 4.0 and later. See [Empty Relations](compatibility/operators.md#empty-relations) for native-input support and writer fallback. |
-| `InMemoryTableScanExec` | ⚠️ | Experimental, disabled by default. Set `spark.comet.exec.inMemoryCache.enabled=true` before the application starts so Comet installs its Arrow cache serializer. Relations with unsupported column types stay in Spark's cache format and fall back. See [In-Memory Cache](in-memory-cache.md). |
+| Operator | Status | Notes |
+| ----------------------- | ------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
+| `FileSourceScanExec` | ✅ | Parquet only. Some types and configurations fall back. See [Parquet Scan Compatibility](compatibility/scans.md). |
+| `BatchScanExec` | ✅ | Apache Iceberg Parquet scans run natively. Native CSV scans are experimental and disabled by default. DataSource V2 Parquet scans are not accelerated. See [Parquet Scan Compatibility](compatibility/scans.md) and the [Iceberg Guide](iceberg.md). |
+| `LocalTableScanExec` | ⚠️ | Disabled by default; there is no acceleration advantage and this operator is typically only used in test code. Can be opted into via config ([#4393](https://github.com/apache/datafusion-comet/pull/4393)). |
+| `EmptyRelationExec` | ✅ | Spark 4.0 and later. See [Empty Relations](compatibility/operators.md#empty-relations) for native-input support and writer fallback. |
+| `RangeExec` | ⚠️ | Disabled by default. Set `spark.comet.exec.range.enabled=true` to generate the rows of `spark.range` and SQL `range()` in native code, so the operators above them run natively. It can be slower than Spark when those operators are only cheap expressions, such as a filter, which Spark compiles together with the range into one loop. Ranges whose arithmetic overflows the `Long` range fall back to Spark. |
+| `InMemoryTableScanExec` | ⚠️ | Experimental, disabled by default. Set `spark.comet.exec.inMemoryCache.enabled=true` before the application starts so Comet installs its Arrow cache serializer. Relations with unsupported column types stay in Spark's cache format and fall back. See [In-Memory Cache](in-memory-cache.md). |
## Projection and filtering
diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md
index bd67aef804..fcbd995a79 100644
--- a/docs/source/user-guide/latest/understanding-comet-plans.md
+++ b/docs/source/user-guide/latest/understanding-comet-plans.md
@@ -292,15 +292,16 @@ by role. Names match what is shown in the plan output.
### Scans
-| Node | Description |
-| ------------------------ | --------------------------------------------------------------------------------------------- |
-| `CometBatchScan` | DataSource V2 scan, including Iceberg Parquet, that produces Arrow batches consumed by Comet. |
-| `CometNativeScan` | Parquet scan that runs entirely in Rust via DataFusion. |
-| `CometIcebergNativeScan` | Iceberg Parquet scan that runs entirely in Rust via DataFusion. |
-| `CometCsvNativeScan` | CSV scan that runs entirely in Rust via DataFusion (experimental). |
-| `CometInMemoryTableScan` | JVM-side scan of a table cached in Comet's Arrow format (experimental, disabled by default). |
-| `CometLocalTableScan` | JVM-side scan of a `LocalTableScanExec` (disabled by default). |
-| `CometEmptyRelation` | Empty native input that replaces `EmptyRelationExec` (Spark 4.0 and later). |
+| Node | Description |
+| ------------------------ | ------------------------------------------------------------------------------------------------------------------ |
+| `CometBatchScan` | DataSource V2 scan, including Iceberg Parquet, that produces Arrow batches consumed by Comet. |
+| `CometNativeScan` | Parquet scan that runs entirely in Rust via DataFusion. |
+| `CometIcebergNativeScan` | Iceberg Parquet scan that runs entirely in Rust via DataFusion. |
+| `CometCsvNativeScan` | CSV scan that runs entirely in Rust via DataFusion (experimental). |
+| `CometInMemoryTableScan` | JVM-side scan of a table cached in Comet's Arrow format (experimental, disabled by default). |
+| `CometLocalTableScan` | JVM-side scan of a `LocalTableScanExec` (disabled by default). |
+| `CometEmptyRelation` | Empty native input that replaces `EmptyRelationExec` (Spark 4.0 and later). |
+| `CometRange` | Generates the rows of `spark.range` and SQL `range()` in native code, replacing `RangeExec` (disabled by default). |
### Native Rust Operators
diff --git a/native/Cargo.lock b/native/Cargo.lock
index e9ec2ff973..07011be0d1 100644
--- a/native/Cargo.lock
+++ b/native/Cargo.lock
@@ -2071,6 +2071,7 @@ dependencies = [
"datafusion-comet-common",
"datafusion-comet-spark-expr",
"futures",
+ "parking_lot",
"tokio",
]
diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs
index 0a0038c5bb..44d2cab14d 100644
--- a/native/core/src/execution/jni_api.rs
+++ b/native/core/src/execution/jni_api.rs
@@ -486,6 +486,7 @@ fn op_name(op: &OpStruct) -> &'static str {
OpStruct::ShuffleScan(_) => "ShuffleScan",
OpStruct::BroadcastNestedLoopJoin(_) => "BroadcastNestedLoopJoin",
OpStruct::Sample(_) => "Sample",
+ OpStruct::RangeScan(_) => "RangeScan",
OpStruct::ContribScan(_) => "ContribScan",
OpStruct::WindowGroupLimit(_) => "WindowGroupLimit",
}
diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs
index 26336affeb..c4d498578b 100644
--- a/native/core/src/execution/planner.rs
+++ b/native/core/src/execution/planner.rs
@@ -81,7 +81,8 @@ use datafusion::{
prelude::SessionContext,
};
use datafusion_comet_operators::{
- CometFilterExec, ExpandExec, ExplodeExec, PartitionedRankLimitExec, SampleExec, WindowFnKind,
+ range_exec, CometFilterExec, ExpandExec, ExplodeExec, PartitionedRankLimitExec, SampleExec,
+ WindowFnKind,
};
use datafusion_comet_spark_expr::{
create_comet_physical_fun, create_comet_physical_fun_with_eval_mode, BinaryOutputStyle,
@@ -1581,6 +1582,22 @@ impl PhysicalPlanner {
Arc::new(SparkPlan::new(spark_plan.plan_id, limit, vec![child])),
))
}
+ OpStruct::RangeScan(range) => {
+ // A leaf with no JVM input: each task produces its own partition of the range.
+ let range: Arc<dyn ExecutionPlan> = Arc::new(range_exec(
+ range.start,
+ range.step,
+ range.num_elements,
+ range.num_slices,
+ self.partition,
+ self.session_ctx.copied_config().batch_size(),
+ )?);
+ Ok((
+ vec![],
+ vec![],
+ Arc::new(SparkPlan::new(spark_plan.plan_id, range, vec![])),
+ ))
+ }
OpStruct::Sample(sample) => {
assert_eq!(children.len(), 1);
let (scans, shuffle_scans, child) =
diff --git a/native/core/src/execution/planner/operator_registry.rs b/native/core/src/execution/planner/operator_registry.rs
index 4f44efdce8..b8f05ba48b 100644
--- a/native/core/src/execution/planner/operator_registry.rs
+++ b/native/core/src/execution/planner/operator_registry.rs
@@ -177,6 +177,7 @@ fn get_operator_type(spark_operator: &Operator) -> Option<OperatorType> {
OpStruct::CsvScan(_) => Some(OperatorType::CsvScan),
OpStruct::ShuffleScan(_) => Some(OperatorType::ShuffleScan),
OpStruct::BroadcastNestedLoopJoin(_) => None,
+ OpStruct::RangeScan(_) => None,
OpStruct::Sample(_) => None, // Not yet in OperatorType enum
// Generic extension point for out-of-tree contrib scans (Delta, Lance, ...); not in
// OperatorType enum. The arm stays unconditional even in non-contrib builds because the
diff --git a/native/operators/Cargo.toml b/native/operators/Cargo.toml
index e43aebe19a..bb45918763 100644
--- a/native/operators/Cargo.toml
+++ b/native/operators/Cargo.toml
@@ -34,6 +34,7 @@ datafusion = { workspace = true }
datafusion-comet-common = { workspace = true }
datafusion-comet-spark-expr = { workspace = true }
futures = { workspace = true }
+parking_lot = "0.12.5"
[dev-dependencies]
criterion = { version = "0.7", features = ["async", "async_tokio", "async_std"] }
diff --git a/native/operators/src/lib.rs b/native/operators/src/lib.rs
index 37075d942c..cfc92a0eb6 100644
--- a/native/operators/src/lib.rs
+++ b/native/operators/src/lib.rs
@@ -28,11 +28,13 @@
mod expand;
mod explode;
mod filter;
+mod range;
mod rank_limit;
mod sample;
pub use expand::ExpandExec;
pub use explode::ExplodeExec;
pub use filter::CometFilterExec;
+pub use range::range_exec;
pub use rank_limit::{PartitionedRankLimitExec, WindowFnKind};
pub use sample::SampleExec;
diff --git a/native/operators/src/range.rs b/native/operators/src/range.rs
new file mode 100644
index 0000000000..3766c82b1a
--- /dev/null
+++ b/native/operators/src/range.rs
@@ -0,0 +1,260 @@
+// 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.
+
+use arrow::array::{Int64Array, RecordBatch};
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use datafusion::common::Result;
+use datafusion::physical_plan::memory::{LazyBatchGenerator, LazyMemoryExec};
+use datafusion::physical_plan::Partitioning;
+use parking_lot::RwLock;
+use std::any::Any;
+use std::fmt;
+use std::sync::Arc;
+
+/// Builds the plan for one Spark partition of `spark.range` or SQL `range()`: a single non-null
+/// `BIGINT` column, produced `batch_size` rows at a time.
+///
+/// Partition bounds follow Spark's `RangeExec`. The JVM side sends only ranges whose arithmetic
+/// cannot overflow in Spark, where both Spark's generated code and its interpreted path produce
+/// the values `partition_start + i * step`. Those values fit in a long, so computing them with
+/// wrapping arithmetic gives them exactly even where `i * step` alone does not fit.
+///
+/// `num_elements` is Spark's element count, and `partition` is the Spark partition index.
+pub fn range_exec(
+ start: i64,
+ step: i64,
+ num_elements: i64,
+ num_slices: i32,
+ partition: i32,
+ batch_size: usize,
+) -> Result<LazyMemoryExec> {
+ let (partition_start, partition_elements) =
+ partition_bounds(start, step, num_elements, num_slices, partition);
+ let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
+ let generator = RangeGenerator {
+ schema: Arc::clone(&schema),
+ step,
+ partition_start,
+ partition_elements,
+ batch_size,
+ produced: 0,
+ };
+ let mut exec = LazyMemoryExec::try_new(schema, vec![Arc::new(RwLock::new(generator))])?;
+ exec.try_set_partitioning(Partitioning::UnknownPartitioning(1))?;
+ Ok(exec)
+}
+
+/// The first value of partition `partition` and the number of values in it, computed the way
+/// Spark's `RangeExec` does (`initRange` in its generated code): in `BigInteger` arithmetic, with
+/// the partition's start and end clamped to the `Long` range.
+///
+/// `i128` holds every intermediate value. The product `partition * num_elements` is below 2^94,
+/// and dividing by `num_slices` before multiplying by `step` keeps the next product below 2^126.
+/// Division truncates toward zero and the remainder takes the dividend's sign, as with
+/// `BigInteger`.
+fn partition_bounds(
+ start: i64,
+ step: i64,
+ num_elements: i64,
+ num_slices: i32,
+ partition: i32,
+) -> (i64, i64) {
+ fn clamp(value: i128) -> i64 {
+ value.clamp(i64::MIN as i128, i64::MAX as i128) as i64
+ }
+
+ let index = partition as i128;
+ let start = start as i128;
+ let step = step as i128;
+ let num_slices = num_slices as i128;
+ let num_elements = num_elements as i128;
+ let partition_start = clamp(index * num_elements / num_slices * step + start);
+ let partition_end = clamp((index + 1) * num_elements / num_slices * step + start);
+ let start_to_end = partition_end as i128 - partition_start as i128;
+ let count = (start_to_end / step) as i64;
+ if count < 0 {
+ (partition_start, 0)
+ } else if start_to_end % step != 0 {
+ (partition_start, count + 1)
+ } else {
+ (partition_start, count)
+ }
+}
+
+/// Generates one partition's values.
+#[derive(Debug)]
+struct RangeGenerator {
+ schema: SchemaRef,
+ step: i64,
+ partition_start: i64,
+ partition_elements: i64,
+ batch_size: usize,
+ /// The number of values produced so far.
+ produced: i64,
+}
+
+impl fmt::Display for RangeGenerator {
+ fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+ write!(
+ f,
+ "CometRange: partition_start={}, step={}, partition_elements={}",
+ self.partition_start, self.step, self.partition_elements
+ )
+ }
+}
+
+impl LazyBatchGenerator for RangeGenerator {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn generate_next_batch(&mut self) -> Result<Option<RecordBatch>> {
+ let first = self.produced;
+ let count = (self.partition_elements - first).min(self.batch_size as i64);
+ if count <= 0 {
+ return Ok(None);
+ }
+ self.produced += count;
+ let (start, step) = (self.partition_start, self.step);
+ let values = Int64Array::from_iter_values(
+ (first..first + count).map(|i| start.wrapping_add(i.wrapping_mul(step))),
+ );
+ Ok(Some(RecordBatch::try_new(
+ Arc::clone(&self.schema),
+ vec![Arc::new(values)],
+ )?))
+ }
+
+ fn reset_state(&self) -> Arc<RwLock<dyn LazyBatchGenerator>> {
+ Arc::new(RwLock::new(RangeGenerator {
+ schema: Arc::clone(&self.schema),
+ produced: 0,
+ ..*self
+ }))
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use datafusion::physical_plan::ExecutionPlan;
+ use datafusion::prelude::SessionContext;
+ use futures::StreamExt;
+
+ async fn collect(
+ start: i64,
+ step: i64,
+ num_elements: i64,
+ num_slices: i32,
+ batch_size: usize,
+ ) -> Vec<i64> {
+ let ctx = SessionContext::new();
+ let mut values = Vec::new();
+ for partition in 0..num_slices {
+ let range =
+ range_exec(start, step, num_elements, num_slices, partition, batch_size).unwrap();
+ let mut stream = range.execute(0, ctx.task_ctx()).unwrap();
+ while let Some(batch) = stream.next().await {
+ let batch = batch.unwrap();
+ assert!(batch.num_rows() <= batch_size);
+ let column = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ values.extend(column.values().iter().copied());
+ }
+ }
+ values
+ }
+
+ /// Spark's element count for `range(start, end, step)`, which these tests only call where it
+ /// fits in a long.
+ fn num_elements(start: i64, end: i64, step: i64) -> i64 {
+ let (start, end, step) = (start as i128, end as i128, step as i128);
+ let n = if (end - start) % step == 0 || (end > start) != (step > 0) {
+ (end - start) / step
+ } else {
+ (end - start) / step + 1
+ };
+ n as i64
+ }
+
+ #[tokio::test]
+ async fn test_values_across_partitions_and_batches() {
+ for (start, end, step, slices) in [
+ (0i64, 10i64, 1i64, 3),
+ (0, 10, 1, 15),
+ (3, 15, 3, 2),
+ (100, -100, -2, 3),
+ (-5000, 5000, 7, 4),
+ ] {
+ let expected: Vec<i64> = if step > 0 {
+ (start..end).step_by(step as usize).collect()
+ } else {
+ (end + 1..=start).rev().step_by((-step) as usize).collect()
+ };
+ for batch_size in [1, 7, 1000, 8192] {
+ let n = num_elements(start, end, step);
+ assert_eq!(
+ collect(start, step, n, slices, batch_size).await,
+ expected,
+ "range({start}, {end}, {step}, {slices}) with {batch_size}-row batches"
+ );
+ }
+ }
+ }
+
+ /// The cases from Spark's `DataFrameRangeSuite` whose bounds are at the edges of the `Long`
+ /// range.
+ #[tokio::test]
+ async fn test_long_range_edges() {
+ let n = num_elements(i64::MIN, i64::MAX, i64::MAX);
+ assert_eq!(
+ collect(i64::MIN, i64::MAX, n, 100, 8192).await,
+ vec![i64::MIN, -1, i64::MAX - 1]
+ );
+ let n = num_elements(i64::MAX, i64::MIN, i64::MIN);
+ assert_eq!(
+ collect(i64::MAX, i64::MIN, n, 100, 8192).await,
+ vec![i64::MAX, -1]
+ );
+ }
+
+ /// Spanning nearly the whole `Long` range, `i * step` stops fitting in a long before the
+ /// values do.
+ #[tokio::test]
+ async fn test_values_where_the_offset_does_not_fit() {
+ let step = 1i64 << 52;
+ let n = num_elements(i64::MIN, i64::MAX, step);
+ let expected: Vec<i64> = (0..n as i128)
+ .map(|i| (i64::MIN as i128 + i * step as i128) as i64)
+ .collect();
+ assert_eq!(collect(i64::MIN, step, n, 1, 1000).await, expected);
+ }
+
+ #[test]
+ fn test_partition_bounds() {
+ // range(0, 10, 1, 3): Spark puts 3, 3 and 4 values in the three partitions.
+ assert_eq!(partition_bounds(0, 1, 10, 3, 0), (0, 3));
+ assert_eq!(partition_bounds(0, 1, 10, 3, 1), (3, 3));
+ assert_eq!(partition_bounds(0, 1, 10, 3, 2), (6, 4));
+ // More slices than values leaves some partitions empty.
+ assert_eq!(partition_bounds(0, 1, 10, 15, 0), (0, 0));
+ assert_eq!(partition_bounds(0, 1, 10, 15, 14), (9, 1));
+ }
+}
diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto
index d6676ac336..1901b59655 100644
--- a/native/proto/src/proto/operator.proto
+++ b/native/proto/src/proto/operator.proto
@@ -68,6 +68,7 @@ message Operator {
Sample sample = 118;
WindowGroupLimit window_group_limit = 119;
IcebergWrite iceberg_write = 120;
+ RangeScan range_scan = 121;
// Extension point for optional, out-of-tree contrib scans (Delta, Lance, ...). The concrete
// scan message (e.g. `DeltaScan`) is packed into this envelope on the JVM side and dispatched
// by `type_url` on the native side. Using a single permanent field -- rather than a new oneof
@@ -843,6 +844,17 @@ message Sample {
int64 seed = 3;
}
+// Spark's `RangeExec`, which produces the rows of `spark.range` and SQL `range()`. The native side
+// computes its partition's values from these fields and the partition index, the way Spark's
+// `RangeExec` does.
+message RangeScan {
+ int64 start = 1;
+ int64 step = 2;
+ // Spark's element count. Comet sends only ranges whose count fits in a long.
+ int64 num_elements = 3;
+ int32 num_slices = 4;
+}
+
enum CompressionCodec {
None = 0;
Zstd = 1;
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index 11ae8ac041..e15b32d6e9 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -264,6 +264,15 @@ object CometConf extends ShimCometConf {
createExecEnabledConfig("takeOrderedAndProject", defaultValue = true)
val COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED: ConfigEntry[Boolean] =
createExecEnabledConfig("localTableScan", defaultValue = false)
+ val COMET_EXEC_RANGE_ENABLED: ConfigEntry[Boolean] =
+ createExecEnabledConfig(
+ "range",
+ defaultValue = false,
+ notes = Some(
+ "When enabled, Comet generates the rows of `spark.range` and SQL `range()` in native " +
+ "code, so the operators above them can run natively. It is off by default because " +
+ "it can be slower than Spark when those operators are only cheap expressions, such " +
+ "as a filter, which Spark compiles together with the range into one loop"))
val COMET_EXEC_EMPTY_RELATION_ENABLED: ConfigEntry[Boolean] =
createExecEnabledConfig("emptyRelation", defaultValue = true)
val COMET_EXEC_SAMPLE_ENABLED: ConfigEntry[Boolean] =
diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
index ce354aed33..e5e250f75d 100644
--- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
@@ -107,6 +107,7 @@ object CometExecRule {
classOf[SortMergeJoinExec] -> CometSortMergeJoinExec,
classOf[SortExec] -> CometSortExec,
classOf[LocalTableScanExec] -> CometLocalTableScanExec,
+ classOf[RangeExec] -> CometRangeExec,
classOf[InMemoryTableScanExec] -> CometInMemoryTableScanExec,
classOf[SampleExec] -> CometSampleExec,
classOf[WindowExec] -> CometWindowExec) ++
@@ -447,6 +448,19 @@ case class CometExecRule(session: SparkSession)
case c: CometSparkToColumnarExec =>
convertToComet(c, CometScanWrapper).getOrElse(c)
+ // A leaf with its own enabled Comet operator, such as RangeExec, uses that operator. The
+ // Spark-to-Arrow conversion is the fallback for a leaf the operator declines.
+ case op: LeafExecNode if hasEnabledHandler(op) =>
+ convertToComet(op, allExecs(op.getClass))
+ .orElse {
+ if (shouldApplySparkToColumnar(conf, op)) {
+ convertToComet(op, CometSparkToColumnarExec)
+ } else {
+ None
+ }
+ }
+ .getOrElse(op)
+
case op if shouldApplySparkToColumnar(conf, op) =>
convertToComet(op, CometSparkToColumnarExec).getOrElse(op)
@@ -1149,6 +1163,9 @@ case class CometExecRule(session: SparkSession)
}
}
+ private def hasEnabledHandler(op: SparkPlan): Boolean =
+ allExecs.get(op.getClass).exists(_.enabledConfig.forall(_.get(op.conf)))
+
private def shouldApplySparkToColumnar(conf: SQLConf, op: SparkPlan): Boolean = {
// Only consider converting leaf nodes to columnar currently, so that all the following
// operators can have a chance to be converted to columnar. Leaf operators that output
diff --git a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala
index b576961e46..dcaa22c957 100644
--- a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala
@@ -24,7 +24,7 @@ import java.util.IdentityHashMap
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.util.sideBySide
-import org.apache.spark.sql.comet.{CometCollectLimitExec, CometColumnarToRowExec, CometIcebergWriteExec, CometMapInBatchExec, CometNativeColumnarToRowExec, CometNativeWriteExec, CometPlan, CometSparkToColumnarExec}
+import org.apache.spark.sql.comet.{CometCollectLimitExec, CometColumnarToRowExec, CometIcebergWriteExec, CometMapInBatchExec, CometNativeColumnarToRowExec, CometNativeWriteExec, CometPlan, CometRangeExec, CometSparkToColumnarExec}
import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometShuffleExchangeExec}
import org.apache.spark.sql.comet.shims.{MapInBatchInfo, ShimCometMapInBatch}
import org.apache.spark.sql.comet.util.Utils.containsVariantType
@@ -107,6 +107,10 @@ case class EliminateRedundantTransitions(session: SparkSession)
// Write should be final operation in the plan
case ColumnarToRowExec(nativeWrite: CometNativeWriteExec) =>
nativeWrite
+ // A range that no native operator consumes would be generated natively only to be turned
+ // back into rows, so keep Spark's RangeExec, which whole-stage codegen fuses with its parent.
+ case ColumnarToRowExec(range: CometRangeExec) =>
+ range.originalPlan
case c @ ColumnarToRowExec(child) if hasCometNativeChild(child, containsCometPlanMemo) =>
val op = createColumnarToRowExec(child)
if (c.logicalLink.isEmpty) {
@@ -120,6 +124,10 @@ case class EliminateRedundantTransitions(session: SparkSession)
sparkToColumnar.child
case CometNativeColumnarToRowExec(sparkToColumnar: CometSparkToColumnarExec) =>
sparkToColumnar.child
+ case CometColumnarToRowExec(range: CometRangeExec) =>
+ range.originalPlan
+ case CometNativeColumnarToRowExec(range: CometRangeExec) =>
+ range.originalPlan
case CometSparkToColumnarExec(child: CometSparkToColumnarExec) => child
// Replace MapInBatchExec (PythonMapInArrowExec / MapInArrowExec / MapInPandasExec) that has
// a ColumnarToRow child with CometMapInBatchExec, eliminating the input and output
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometRangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometRangeExec.scala
new file mode 100644
index 0000000000..3cb049398e
--- /dev/null
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometRangeExec.scala
@@ -0,0 +1,127 @@
+/*
+ * 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 org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.execution.{RangeExec, SparkPlan}
+
+import com.google.common.base.Objects
+
+import org.apache.comet.{CometConf, ConfigEntry}
+import org.apache.comet.serde.{CometOperatorSerde, Compatible, OperatorOuterClass, SupportLevel, Unsupported}
+import org.apache.comet.serde.OperatorOuterClass.Operator
+
+/**
+ * Comet's version of Spark's `RangeExec`, which produces the rows of `spark.range` and SQL
+ * `range()`. The values are generated in native code (`range_exec` in the native operators
+ * crate). Each task computes its own partition from the partition index, so this operator reports
+ * Spark's partitioning, and the native plan it belongs to runs one task per slice with no JVM
+ * input. Partitions, values and their order match Spark's `RangeExec`. Where no native operator
+ * consumes the range, `EliminateRedundantTransitions` restores Spark's `RangeExec`.
+ */
+case class CometRangeExec(
+ override val nativeOp: Operator,
+ @transient override val originalPlan: RangeExec,
+ override val output: Seq[Attribute],
+ override val serializedPlanOpt: SerializedPlan)
+ extends CometLeafExec {
+
+ override def simpleString(maxFields: Int): String = {
+ s"$nodeName (${originalPlan.start}, ${originalPlan.end}, step=${originalPlan.step}, " +
+ s"splits=${originalPlan.numSlices})"
+ }
+
+ override protected def doCanonicalize(): SparkPlan = {
+ val canonical = originalPlan.canonicalized.asInstanceOf[RangeExec]
+ CometRangeExec(nativeOp, canonical, canonical.output, SerializedPlan(None))
+ }
+
+ // `originalPlan` carries every parameter that decides the rows: start, end, step and the
+ // number of slices.
+ override def equals(obj: Any): Boolean = {
+ obj match {
+ case other: CometRangeExec =>
+ this.originalPlan == other.originalPlan &&
+ this.output == other.output &&
+ this.serializedPlanOpt == other.serializedPlanOpt
+ case _ =>
+ false
+ }
+ }
+
+ override def hashCode(): Int = Objects.hashCode(originalPlan, output)
+}
+
+object CometRangeExec extends CometOperatorSerde[RangeExec] {
+
+ override def enabledConfig: Option[ConfigEntry[Boolean]] = Some(
+ CometConf.COMET_EXEC_RANGE_ENABLED)
+
+ /**
+ * A range whose arithmetic may overflow stays on Spark. Spark's generated code for `RangeExec`
+ * and its interpreted `RangeExec.doExecute`, which runs when whole-stage codegen is off or the
+ * generated code is too large, can return different rows for it. So does a range with fewer
+ * than one slice, which Spark fails when it runs and Comet would return empty.
+ */
+ override def getSupportLevel(op: RangeExec): SupportLevel = {
+ if (op.isEmptyRange) {
+ // Spark produces no partitions for it whatever its slice count, and so does Comet.
+ Compatible()
+ } else if (op.numSlices < 1) {
+ Unsupported(Some(s"Spark fails a range with ${op.numSlices} slices"))
+ } else if (mayOverflow(op)) {
+ Unsupported(
+ Some(
+ "Spark can return different rows for a range whose arithmetic overflows, " +
+ "depending on whether whole-stage codegen runs it"))
+ } else {
+ Compatible()
+ }
+ }
+
+ /**
+ * Whether Spark's generated code for `RangeExec` can overflow for `op`. It reads the element
+ * count as a long, and walks each slice in batches of up to 1000 values, computing a batch's
+ * end as its start plus the batch's span. The span wraps if it does not fit in a long. No slice
+ * holds more than `ceil(numElements / numSlices)` values. Otherwise both of Spark's paths
+ * produce the values `start + i * step` that Comet does.
+ */
+ private def mayOverflow(op: RangeExec): Boolean = {
+ val maxSliceElements = (op.numElements + op.numSlices - 1) / op.numSlices
+ !op.numElements.isValidLong || !(BigInt(op.step) * maxSliceElements.min(1000)).isValidLong
+ }
+
+ override def convert(
+ op: RangeExec,
+ builder: Operator.Builder,
+ childOp: Operator*): Option[Operator] = {
+ val rangeScan = OperatorOuterClass.RangeScan
+ .newBuilder()
+ .setStart(op.start)
+ .setStep(op.step)
+ // getSupportLevel declines a count that does not fit in a long.
+ .setNumElements(op.numElements.toLong)
+ .setNumSlices(op.numSlices)
+ Some(builder.setRangeScan(rangeScan).build())
+ }
+
+ override def createExec(nativeOp: Operator, op: RangeExec): CometNativeExec =
+ CometRangeExec(nativeOp, op, op.output, SerializedPlan(None))
+}
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometRangeExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometRangeExecSuite.scala
new file mode 100644
index 0000000000..ba9f5a68d7
--- /dev/null
+++ b/spark/src/test/scala/org/apache/comet/exec/CometRangeExecSuite.scala
@@ -0,0 +1,233 @@
+/*
+ * 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 scala.util.Random
+
+import org.apache.spark.SparkConf
+import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
+import org.apache.spark.sql.comet.{CometRangeExec, CometSparkToColumnarExec}
+import org.apache.spark.sql.execution.{RangeExec, SparkPlan}
+import org.apache.spark.sql.functions.{col, count, sum}
+import org.apache.spark.sql.internal.SQLConf
+
+import org.apache.comet.CometConf
+
+class CometRangeExecSuite extends CometTestBase {
+
+ // CometTestBase enables the Spark-to-Arrow conversion, which also applies to Range. Use its
+ // production default so these tests see CometRangeExec on its own.
+ override protected def sparkConf: SparkConf =
+ super.sparkConf.remove(CometConf.COMET_SPARK_TO_ARROW_ENABLED.key)
+
+ /** Defines a test that runs with `CometRangeExec` enabled. */
+ private def rangeTest(name: String)(f: => Unit): Unit =
+ test(name) {
+ withSQLConf(CometConf.COMET_EXEC_RANGE_ENABLED.key -> "true")(f)
+ }
+
+ private def cometRanges(plan: SparkPlan): Seq[CometRangeExec] =
+ collect(plan) { case r: CometRangeExec => r }
+
+ private def checkCometRange(df: => DataFrame): SparkPlan = {
+ val (_, cometPlan) = checkSparkAnswerAndOperator(df, Seq(classOf[CometRangeExec]))
+ cometPlan
+ }
+
+ // (start, end, step, splits). The first rows are the cases in Spark's DataFrameRangeSuite.
+ private val ranges: Seq[(Long, Long, Long, Int)] = Seq(
+ (0L, 10L, 1L, 15),
+ (3L, 15L, 3L, 2),
+ (1L, -2L, -2L, 6),
+ (-3L, -8L, -2L, 1),
+ (-8L, -4L, 2L, 1),
+ (Long.MinValue, Long.MaxValue, Long.MaxValue, 100),
+ (Long.MaxValue, Long.MinValue, Long.MinValue, 100),
+ (-1000000L, 1000000L, 111111L, 4),
+ (0L, 100L, 2L, 3),
+ (100L, -100L, -2L, 3),
+ (-1500L, 1500L, 3L, 5),
+ (10L, 0L, -1L, 1),
+ // More slices than values, and a step larger than the range.
+ (0L, 3L, 1L, 8),
+ (0L, 10L, 100L, 4),
+ // Ranges that end at the edges of the Long range.
+ (Long.MaxValue - 10, Long.MaxValue, 3L, 2),
+ (Long.MinValue + 10, Long.MinValue, -3L, 2),
+ // A product of the index and the step that does not fit in a long, with values that do.
+ (Long.MinValue, Long.MaxValue, 1L << 52, 1))
+
+ // Ranges with no rows, which Spark plans as RangeExec over an empty RDD.
+ private val emptyRanges: Seq[(Long, Long, Long, Int)] = Seq(
+ (1L, -2L, 1L, 4),
+ (-10L, -9L, -20L, 1),
+ (Long.MaxValue - 3, Long.MinValue + 2, 1L, 2),
+ (Long.MaxValue - 3, Long.MaxValue - 3, 1L, 2))
+
+ rangeTest("range is replaced and feeds the native operators above it") {
+ val plan = checkCometRange(
+ spark.range(0, 1000, 3, 4).selectExpr("id", "id * 2 AS doubled").filter("id % 2 = 0"))
+ val range = cometRanges(plan).head
+ assert(range.outputPartitioning == range.originalPlan.outputPartitioning)
+ assert(range.outputOrdering == range.originalPlan.outputOrdering)
+ }
+
+ rangeTest("range matches Spark for edge-case bounds, steps and splits") {
+ ranges.foreach { case (start, end, step, splits) =>
+ withClue(s"range($start, $end, $step, $splits): ") {
+ checkCometRange(spark.range(start, end, step, splits).selectExpr("id", "id % 7 AS m"))
+ }
+ }
+ }
+
+ rangeTest("a range no native operator consumes stays on Spark") {
+ // Generating it natively would only add a conversion back to rows.
+ val df = spark.range(0, 100, 1, 2).toDF()
+ val (_, plan) = checkSparkAnswer(df)
+ assert(cometRanges(plan).isEmpty, plan)
+ assert(collect(plan) { case r: RangeExec => r }.nonEmpty, plan)
+ }
+
+ rangeTest("empty ranges") {
+ emptyRanges.foreach { case (start, end, step, splits) =>
+ withClue(s"range($start, $end, $step, $splits): ") {
+ val df = spark.range(start, end, step, splits).selectExpr("id + 1 AS next")
+ checkCometRange(df)
+ assert(df.collect().isEmpty)
+ }
+ }
+ }
+
+ rangeTest("empty range under a global aggregate") {
+ // The native plan has no partitions, and the aggregate still returns its one row.
+ val df = spark.range(1, -2, 1, 4).agg(count("id"), sum("id"))
+ checkSparkAnswer(df)
+ assert(df.collect().toSeq == Seq(Row(0L, null)))
+ }
+
+ rangeTest("a range whose arithmetic may overflow stays on Spark") {
+ val reason = "Spark can return different rows for a range whose arithmetic overflows"
+ // Spark's generated code returns no rows here, because the end of its batch wraps, and its
+ // interpreted RangeExec returns four.
+ val overflowing = spark.range(Long.MinValue, Long.MaxValue, 1L << 62, 1).selectExpr("id + 1")
+ Seq("true", "false").foreach { codegen =>
+ withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> codegen) {
+ checkSparkAnswerAndFallbackReason(overflowing, reason)
+ }
+ }
+ // An element count that does not fit in a long. Only with codegen, since Spark's interpreted
+ // RangeExec would produce every value.
+ checkSparkAnswerAndFallbackReason(
+ spark.range(Long.MinValue, Long.MaxValue, 1, 4).selectExpr("id + 1"),
+ reason)
+ // A range CometRangeExec declines can still use the Spark-to-Arrow conversion.
+ withSQLConf(CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "true") {
+ val (_, plan) = checkSparkAnswer(overflowing)
+ assert(collect(plan) { case c: CometSparkToColumnarExec => c }.nonEmpty, plan)
+ }
+ }
+
+ rangeTest("a range with fewer than one slice fails as in Spark") {
+ Seq(0, -1).foreach { splits =>
+ withClue(s"$splits slices: ") {
+ val df = spark.range(0, 10, 1, splits).selectExpr("id + 1")
+ val (sparkError, cometError) = checkSparkAnswerMaybeThrows(df)
+ assert(sparkError.isDefined && cometError.isDefined, (sparkError, cometError))
+ assert(cometError.get.getMessage == sparkError.get.getMessage)
+ assert(cometRanges(df.queryExecution.executedPlan).isEmpty)
+ }
+ }
+ }
+
+ rangeTest("range with randomized parameters") {
+ // Mirrors Spark's DataFrameRangeSuite test of the same name.
+ val maxNumSteps = 10L * 1000
+ val seed = System.currentTimeMillis()
+ val random = new Random(seed)
+
+ def randomBound(): Long = {
+ val n = random.nextLong() % (Long.MaxValue / (100 * maxNumSteps))
+ if (random.nextBoolean()) n else -n
+ }
+
+ for (_ <- 1 to 10) {
+ val start = randomBound()
+ val end = randomBound()
+ val numSteps = (math.abs(random.nextLong()) % maxNumSteps) + 1
+ val stepAbs = (math.abs(end - start) / numSteps) + 1
+ val step = if (start < end) stepAbs else -stepAbs
+ val partitions = random.nextInt(20) + 1
+
+ withClue(s"seed = $seed start = $start end = $end step = $step partitions = $partitions") {
+ checkCometRange(spark.range(start, end, step, partitions).agg(count("id"), sum("id")))
+ }
+ }
+ }
+
+ rangeTest("values come out in order across batch boundaries") {
+ // Batches of 7 rows split Spark's 1000-value batches at many different offsets.
+ withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "7") {
+ Seq((0L, 10000L, 3L, 1), (10000L, 0L, -3L, 1), (-5000L, 5000L, 1L, 7)).foreach {
+ case (start, end, step, splits) =>
+ withClue(s"range($start, $end, $step, $splits): ") {
+ val df = spark.range(start, end, step, splits).selectExpr("id", "id * 3 AS tripled")
+ checkCometRange(df)
+ assert(df.collect().map(_.getLong(0)).toSeq == (start until end by step))
+ }
+ }
+ }
+ }
+
+ rangeTest("SQL range() is replaced") {
+ checkCometRange(sql("SELECT id, id * 3 AS tripled FROM range(3)"))
+ checkCometRange(sql("SELECT sum(id), count(*) FROM range(5, 0, -1, 2)"))
+ }
+
+ rangeTest("output rows metric") {
+ val df = spark.range(0, 1000, 1, 3).selectExpr("id + 1")
+ df.collect()
+ val range = cometRanges(df.queryExecution.executedPlan).head
+ // The native operator reports DataFusion's baseline metrics.
+ assert(range.metrics("output_rows").value == 1000)
+ }
+
+ test("range stays on Spark by default") {
+ val df = spark.range(0, 100, 1, 2).selectExpr("id + 1")
+ df.collect()
+ val plan = df.queryExecution.executedPlan
+ assert(cometRanges(plan).isEmpty, plan)
+ assert(collect(plan) { case r: RangeExec => r }.nonEmpty, plan)
+ }
+
+ rangeTest("takes precedence over the Spark-to-Arrow conversion") {
+ withSQLConf(CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "true") {
+ val plan = checkCometRange(spark.range(0, 100, 1, 2).selectExpr("id + 1"))
+ assert(collect(plan) { case c: CometSparkToColumnarExec => c }.isEmpty, plan)
+ }
+ }
+
+ rangeTest("exchanges are reused over equal ranges but not different ones") {
+ def counts(end: Long): DataFrame =
+ spark.range(0, end, 1, 4).groupBy((col("id") % 10).as("k")).count()
+ checkCometRange(counts(100).union(counts(200)))
+ val plan = checkCometRange(counts(100).union(counts(100)))
+ assertExchangeReuseOver(plan, "equal ranges") { case r: CometRangeExec => r }
+ }
+}
diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometRangeBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometRangeBenchmark.scala
new file mode 100644
index 0000000000..7d55e90f02
--- /dev/null
+++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometRangeBenchmark.scala
@@ -0,0 +1,133 @@
+/*
+ * 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.benchmark
+
+import org.apache.spark.benchmark.Benchmark
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.comet.{CometHashAggregateExec, CometPlan, CometRangeExec, CometSparkToColumnarExec}
+import org.apache.spark.sql.execution.{RangeExec, SparkPlan}
+import org.apache.spark.sql.functions.{col, sum}
+
+import org.apache.comet.CometConf
+
+/**
+ * Compares three ways to run a query over `spark.range`:
+ *
+ * - Spark: Comet disabled.
+ * - SparkToColumnar: Spark's `RangeExec`, whose rows `CometSparkToColumnarExec` converts to
+ * Arrow for the native operators above it.
+ * - CometRange: `CometRangeExec`, which generates the values in native code.
+ *
+ * The cases sweep how much work sits above the range, from a sum that Spark's whole-stage codegen
+ * fuses with the range loop to a high-cardinality aggregate. Every arm's result and plan are
+ * checked before it is timed, and the Spark arm runs again at the end of each case to show the
+ * noise. To run this benchmark:
+ * {{{
+ * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometRangeBenchmark
+ * }}}
+ * Results will be written to "spark/benchmarks/CometRangeBenchmark-**results.txt".
+ */
+object CometRangeBenchmark extends CometBenchmarkBase {
+
+ private val numRows = 64L * 1024 * 1024
+
+ private case class Arm(name: String, confs: Seq[(String, String)], check: SparkPlan => Unit)
+
+ private val cometConfs =
+ Seq(CometConf.COMET_ENABLED.key -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true")
+
+ private val arms = Seq(
+ Arm(
+ "Spark",
+ Seq(CometConf.COMET_ENABLED.key -> "false"),
+ plan => {
+ assert(collect(plan) { case r: RangeExec => r }.nonEmpty, plan)
+ assert(collect(plan) { case c: CometPlan => c }.isEmpty, plan)
+ }),
+ Arm(
+ "SparkToColumnar",
+ cometConfs ++ Seq(
+ CometConf.COMET_EXEC_RANGE_ENABLED.key -> "false",
+ CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "true"),
+ plan => {
+ assert(collect(plan) { case c: CometSparkToColumnarExec => c }.nonEmpty, plan)
+ assert(collect(plan) { case a: CometHashAggregateExec => a }.nonEmpty, plan)
+ }),
+ Arm(
+ "CometRange",
+ cometConfs ++ Seq(
+ CometConf.COMET_EXEC_RANGE_ENABLED.key -> "true",
+ CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "false"),
+ plan => {
+ assert(collect(plan) { case r: CometRangeExec => r }.nonEmpty, plan)
+ assert(collect(plan) { case a: CometHashAggregateExec => a }.nonEmpty, plan)
+ }))
+
+ private def groupedCount(keys: Long): DataFrame =
+ spark
+ .range(numRows)
+ .groupBy((col("id") % keys).as("k"))
+ .count()
+ .agg(sum("count"))
+
+ private val cases: Seq[(String, () => DataFrame)] = Seq(
+ "range -> sum" -> (() => spark.range(numRows).selectExpr("sum(id)")),
+ "range -> filter -> sum" -> (() =>
+ spark.range(numRows).filter("id % 3 = 0").selectExpr("sum(id)")),
+ "range -> xxhash64 -> sum" -> (() => spark.range(numRows).selectExpr("sum(xxhash64(id))")),
+ "range -> cast to string -> sum(length)" -> (() =>
+ spark.range(numRows).selectExpr("sum(length(CAST(id AS STRING)))")),
+ "range -> group by 100 keys" -> (() => groupedCount(100)),
+ "range -> group by 1M keys" -> (() => groupedCount(1000000)))
+
+ private def runArm(arm: Arm, query: () => DataFrame): (Seq[Row], SparkPlan) = {
+ var result: (Seq[Row], SparkPlan) = null
+ withSQLConf(arm.confs: _*) {
+ val df = query()
+ result = (df.collect().toSeq, df.queryExecution.executedPlan)
+ }
+ result
+ }
+
+ override def runCometBenchmark(mainArgs: Array[String]): Unit = {
+ cases.foreach { case (name, query) =>
+ runBenchmark(name) {
+ // Check every arm's plan and result before timing anything.
+ val results = arms.map(runArm(_, query))
+ arms.zip(results).foreach { case (arm, (rows, plan)) =>
+ arm.check(plan)
+ assert(
+ rows == results.head._1,
+ s"${arm.name} returned $rows, expected ${results.head._1}")
+ }
+
+ val benchmark = new Benchmark(name, numRows, output = output)
+ (arms :+ arms.head.copy(name = "Spark (repeat)")).foreach { arm =>
+ benchmark.addCase(arm.name) { _ =>
+ withSQLConf(arm.confs: _*) {
+ query().collect()
+ }
+ }
+ }
+ benchmark.run()
+ }
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org