diff --git a/.ai/skills/review-comet-expression-pr/SKILL.md b/.ai/skills/review-comet-expression-pr/SKILL.md index 2c83d9d15ba..17f8e2f04ce 100644 --- a/.ai/skills/review-comet-expression-pr/SKILL.md +++ b/.ai/skills/review-comet-expression-pr/SKILL.md @@ -30,11 +30,12 @@ bar, and the output format. This skill only covers what is specific to expressio ## Read the Contributor Guide First -| Doc | What you need from it | -| ---------------------------------------------------------- | ----------------------------------------------------------------------------------------- | -| `docs/source/contributor-guide/adding_a_new_expression.md` | The serde contract, support levels, when to set the return type explicitly, shimming | -| `docs/source/contributor-guide/sql-file-tests.md` | The test framework expression PRs are expected to use, and every directive it supports | -| `docs/source/contributor-guide/optimizing_expressions.md` | The benchmark workflow and the no-regression rule, for PRs that change an existing kernel | +| Doc | What you need from it | +| ---------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `docs/source/contributor-guide/adding_a_new_expression.md` | The serde contract, support levels, when to set the return type explicitly, shimming | +| `docs/source/contributor-guide/sql-file-tests.md` | The test framework expression PRs are expected to use, and every directive it supports | +| `docs/source/contributor-guide/optimizing_expressions.md` | The benchmark workflow and the no-regression rule, for PRs that change an existing kernel | +| `docs/source/contributor-guide/timezones.md` | The timestamp label rule, where the session timezone comes from, and how to test it, for datetime expressions, casts, and anything that returns a timestamp | Hold the diff against these. If the PR does something a guide says to do differently, either the PR is wrong or the guide is out of date, and you need to say which. @@ -77,6 +78,10 @@ Location: `spark/src/main/scala/org/apache/comet/serde/` - [ ] A change that routes a whole class of expressions to the JVM codegen dispatcher lists the new shapes it admits and tests one of each, including decimal results whose scale differs from the declared type and boolean inputs from sliced batches (#6424, #6425) +- [ ] A timezone-aware expression serializes `CometTimeZone.nativeId(expr.timeZoneId)`, the + timezone Spark stamped on it in a form native code can parse, rather than + `SQLConf.get.sessionLocalTimeZone` or the JVM default, and returns + `CometTimeZone.supportLevel` from `getSupportLevel` when `nativeId` gives `None` ### Registration in `QueryPlanSerde.scala` @@ -112,6 +117,13 @@ Location: `native/spark-expr/src/`, registered in `comet_scalar_funcs.rs`. - [ ] Comparisons, grouping and set operations treat `-0.0` and `0.0`, and NaN payloads, as Spark does, including inside arrays and structs. Where a DataFusion kernel treats them differently, the serde falls back (#5507, #5701). +- [ ] A `TimestampType` result is labelled `"UTC"`, in `return_type()` or `data_type()` and in the + arrays the kernel builds. Not the session timezone, and not `None`. +- [ ] Local-time work uses the session timezone from the proto. DataFusion's own datetime functions + take the timezone from the input's label, so wiring one in for a timezone-aware Spark + expression evaluates it in UTC unless the session timezone is passed explicitly. +- [ ] A UTC fast path matches a fixed list of UTC aliases, as `is_utc_timezone` in + `extract_date_part.rs` does, and sends every other timezone ID down the general path Before accepting a hand-written kernel, ask whether the function already exists upstream in DataFusion or the `datafusion-spark` crate. Comet prefers wiring an upstream function over carrying @@ -169,8 +181,10 @@ single file by appending a substring of its name to the suite argument. - [ ] Edge cases tested: empty input, overflow, boundary values, negative values - [ ] Both literal and column arguments tested, in every combination for multi-argument expressions. They take different code paths. -- [ ] Timezone handling tested for timestamp and datetime expressions, including a non-UTC session - timezone and timestamps with and without timezone +- [ ] Timestamp and datetime expressions are tested as "Testing timezone-sensitive code" in + `timezones.md` describes. That means several session timezones through `ConfigMatrix`, + including `Etc/UTC` and a zone with DST, both `TIMESTAMP` and `TIMESTAMP_NTZ` inputs, and a + result that is compared or fed to another expression rather than only projected. - [ ] SQL syntax gated with `MinSparkVersion` when it only parses on newer Spark - [ ] For a function from another project, such as an Iceberg transform, the expected values come from that project's Java implementation, which is what Spark runs, not from iceberg-rust @@ -254,3 +268,5 @@ reference in that doc is the only place they are documented. 5. **Missing `getSupportLevel`**, divergences left undeclared rather than marked `Incompatible` 6. **Version-specific Spark behavior implemented once**, with no shim 7. **Name collides with a DataFusion built-in** and no explicit return type +8. **Timestamp result mislabelled**, with the session timezone or no timezone instead of `"UTC"`. + It only shows once the result is compared or fed to another expression. diff --git a/.ai/skills/review-comet-ffi-pr/SKILL.md b/.ai/skills/review-comet-ffi-pr/SKILL.md index b42ad33d1e3..88d47c1683c 100644 --- a/.ai/skills/review-comet-ffi-pr/SKILL.md +++ b/.ai/skills/review-comet-ffi-pr/SKILL.md @@ -138,6 +138,11 @@ a new subclass needs no new case as long as `getValueVector` returns an Arrow ve `ColumnarBatchArrowReader` decodes dictionary-encoded columns before export, and `reconcileStreamSchema` advertises the value type. A new reader that exports a dictionary would reach `ScanExec` as is, and only the cast in `build_record_batch` would unpack it. +- [ ] **Timestamps cross unconverted.** Both directions pass the raw microseconds. JVM producers + label `TimestampType` with `CometArrowStream.NATIVE_TIMEZONE`, which is `"UTC"`, and a new + producer must use it too rather than the session timezone. A change that shifts values by a + timezone at the boundary gives wrong answers. See "How Comet represents timestamps" in + `docs/source/contributor-guide/timezones.md`. ## 5. Memory Accounting diff --git a/.ai/skills/review-comet-iceberg-write-pr/SKILL.md b/.ai/skills/review-comet-iceberg-write-pr/SKILL.md index 9752f865a5d..a835f45441a 100644 --- a/.ai/skills/review-comet-iceberg-write-pr/SKILL.md +++ b/.ai/skills/review-comet-iceberg-write-pr/SKILL.md @@ -111,6 +111,16 @@ only the latest. `Float.compare` does not ([#6138](https://github.com/apache/datafusion-comet/issues/6138)). Check equality, hashing, ordering and rendering for float, double, timestamp, timestamptz, binary and decimal. +- [ ] **Timestamp partition values are UTC.** Iceberg's `years`, `months`, `days` and `hours` + ignore the session timezone. For date and timestamp sources, `PartitionValueCalculator` + (`iceberg_partition_value.rs`) computes them with the same kernels in + `iceberg_funcs/temporal.rs` that key the sort in front of a clustered write, and those follow + iceberg-java's `DateTimeUtil`, including its pre-1970 rounding. A change that moves one of + these transforms back to iceberg-rust, or that makes the kernels read the column's timezone + label, breaks that agreement. iceberg-rust's `years` and `months` follow the label + ([apache/iceberg-rust#3142](https://github.com/apache/iceberg-rust/issues/3142)). Ask for a + test in a non-UTC session for a change here. "Partition values" in `iceberg-writes.md` lists + where iceberg-rust differs, and `timezones.md` covers how Comet labels timestamps. - [ ] **Partition paths use the Java renderers.** Directory names come from `CometLocationGenerator` / `partition_to_path` in `iceberg_partition_path.rs`, which follow iceberg-java's `partitionToPath`, with `java_float_string` for floats. A PR that calls @@ -142,6 +152,14 @@ Read the ownership table in the contributor guide before reviewing any change ne `SparkWrite.abort` keeps working. - [ ] The user-visible exception type is still what Spark's own write path throws on each Spark version ([#6143](https://github.com/apache/datafusion-comet/issues/6143)). +- [ ] **Every open file and every held-back row is reserved.** Files must be opened through + `MeteredParquetWriterBuilder`, and rows the writer holds outside iceberg-rust (the + `PartitionFeed`s, or anything a new feed holds back) must be added to what `run_write_task` reserves. A + builder that constructs `ParquetWriterBuilder` directly leaves its files out of the task's + memory reservation, so a wide fanout write grows past the pool instead of failing its task. + A storage scheme newly supported for writes needs its entry in + `StorageWrites::for_location`: a file reports its flushed row groups until its storage has + written them out, which a local file does at once and an object store only part by part. ## 5. Plan Shape @@ -179,16 +197,16 @@ Read the ownership table in the contributor guide before reviewing any change ne ## 7. Tests -| Suite | Covers | -| ----------------------------------- | ------------------------------------------------------------------------------------ | -| `CometIcebergWriteActionSuite` | Both layers end to end, parity with iceberg-java, DML, failure cleanup, AQE re-plans | -| `CometIcebergWriteDetectionSuite` | One case per eligibility rule | -| `CometIcebergSystemFunctionSuite` | Native partition transforms keeping partitioned writes native | -| `CometIcebergRewriteActionSuite` | `rewrite_data_files` through the split plan and the native writer | -| `IcebergWriteProtoTranslationSuite` | Property to `IcebergParquetWriteSettings` translation | -| Rust tests in `iceberg_write.rs` | Rolling, fanout order, clustered checks, cleanup guard, manifest round trip | -| `iceberg_partition_path.rs` tests | Partition path rendering against iceberg-java | -| `CometIcebergWriteBenchmark` | Native versus iceberg-java throughput | +| Suite | Covers | +| ----------------------------------- | ----------------------------------------------------------------------------------------------- | +| `CometIcebergWriteActionSuite` | Both layers end to end, parity with iceberg-java, DML, failure cleanup, AQE re-plans | +| `CometIcebergWriteDetectionSuite` | One case per eligibility rule | +| `CometIcebergSystemFunctionSuite` | Native partition transforms keeping partitioned writes native | +| `CometIcebergRewriteActionSuite` | `rewrite_data_files` through the split plan and the native writer | +| `IcebergWriteProtoTranslationSuite` | Property to `IcebergParquetWriteSettings` translation | +| Rust tests in `iceberg_write.rs` | Rolling, fanout order, clustered checks, cleanup guard, manifest round trip, memory reservation | +| `iceberg_partition_path.rs` tests | Partition path rendering against iceberg-java | +| `CometIcebergWriteBenchmark` | Native versus iceberg-java throughput | Ask specifically: diff --git a/.ai/skills/review-comet-pr/SKILL.md b/.ai/skills/review-comet-pr/SKILL.md index 050a60802ce..e93361b69f9 100644 --- a/.ai/skills/review-comet-pr/SKILL.md +++ b/.ai/skills/review-comet-pr/SKILL.md @@ -47,6 +47,11 @@ change. For a new operator, read `docs/source/contributor-guide/adding_a_new_operator.md` alongside this skill. There is no dedicated operator review skill yet. +For a PR that deals with timestamps or the session timezone, read +`docs/source/contributor-guide/timezones.md` as well, whichever areas it touches. Timezone handling +cuts across the serde, the native kernels, the scans and the JVM/native boundary, so no area skill +owns it. "Timestamps and timezones" in step 5 says what to check. + If the PR falls outside all of these, for example build, CI, docs only, or release tooling, this skill on its own is the review. @@ -176,6 +181,41 @@ the minor version is wrong for every earlier patch. CI builds only the newest pa it can't catch that (#6042, #5701). The pull request CI also runs only the default Spark profile, so logic that depends on the Spark version needs the matching `run-spark-*` labels (#6156). +### Timestamps and timezones + +Timezone bugs are easy to miss in review and in tests. A mislabelled timestamp column passes any +test that only projects it, and a result computed in the wrong timezone looks plausible. A PR deals +with timezones if it touches a datetime expression, a cast to or from a timestamp, how a scan reads +timestamps, the timestamp type at the JVM/native boundary, or anything that reads the session +timezone. Searching the diff flags most of these PRs: + +```shell +gh pr diff --repo apache/datafusion-comet | grep -inE 'time_?zone|zoneid|chrono_tz|timestamp(ntz)?type|timestampmicro|timestamp\(' +``` + +For such a PR, hold the diff against "The invariant" and "Guidelines" in +`docs/source/contributor-guide/timezones.md`. Look for these first: + +- A `TimestampType` value labelled with the session timezone, or with no timezone. Inside a native + plan every `TimestampType` value is labelled exactly `"UTC"`, and every `TimestampNTZType` value + has no timezone. Check the declared type and the arrays the code builds, not only the values. +- A timezone taken from the JVM default or the host. An expression uses the `timeZoneId` Spark + stamped on it, passed through `CometTimeZone.nativeId`, not `SQLConf.get.sessionLocalTimeZone`. +- A path gated on a UTC session. `Etc/UTC` is the session default on Ubuntu and Debian images, so + check what the gate does with it, and that the output there is still labelled `"UTC"` rather than + `"Etc/UTC"`. +- A timezone applied to a `TimestampNTZType` value, other than to convert it to `TimestampType`. +- Tests that use a single session timezone, or only project the result. "Testing timezone-sensitive + code" in the same page says what to ask for. + +`timezones.md` also describes specific code: where the `"UTC"` label is set, which serdes serialize +a timezone, how `CometTimeZone` rewrites timezone IDs, what `array_with_timezone` does, how the +scans adapt timestamps, and which expressions go through the codegen dispatcher. A PR that changes +any of these updates the page in the same PR. The page also documents some known limitations as +current behavior, such as chrono-tz's DST horizon and the timezone database versions. A fix for one +of the bugs tracked in [#6335](https://github.com/apache/datafusion-comet/issues/6335) usually +changes that text too. + ### Configuration New configs go in `CometConf.scala` and must follow diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index d6d487eb146..49eafc6e3a1 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -460,6 +460,7 @@ jobs: value: | org.apache.comet.parquet.CometParquetWriterSuite org.apache.comet.parquet.CometEmptyRelationParquetWriterSuite + org.apache.spark.sql.comet.CometNativeWriteSuite org.apache.comet.parquet.ParquetReadV1Suite org.apache.comet.parquet.ParquetTimestampLtzAsNtzSuite org.apache.spark.sql.comet.ParquetEncryptionITCase @@ -486,6 +487,7 @@ jobs: org.apache.comet.exec.CometShuffle4_0Suite org.apache.comet.exec.CometNativeColumnarToRowSuite org.apache.comet.exec.CometNativeShuffleSuite + org.apache.comet.exec.CometShuffleInputConversionSuite org.apache.comet.shuffle.CelebornShufflePartitionPusherSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManagerSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornNativeShuffleWriterSuite @@ -503,6 +505,7 @@ jobs: org.apache.comet.exec.CometShuffleManagerSuite org.apache.comet.exec.DisableAQECometShuffleSuite org.apache.spark.shuffle.comet.CometUnboundedShuffleMemoryAllocatorSuite + org.apache.spark.shuffle.comet.CometUnifiedShuffleMemoryAllocatorSuite org.apache.spark.shuffle.sort.SpillSorterSuite - name: "exec" value: | @@ -563,6 +566,7 @@ jobs: org.apache.spark.sql.comet.PlanDataInjectorShuffleLifecycleSuite org.apache.spark.sql.comet.CometDecimalArithmeticViewSuite org.apache.spark.sql.comet.CometDecimalPromotionSuite + org.apache.spark.sql.comet.DecimalPrecisionSuite org.apache.spark.sql.comet.CometScanWithPlanDataSuite org.apache.spark.sql.comet.util.UtilsSuite org.apache.comet.vector.NativeUtilSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index bf6a86b8f42..22bd6251dcf 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -164,6 +164,7 @@ jobs: value: | org.apache.comet.parquet.CometParquetWriterSuite org.apache.comet.parquet.CometEmptyRelationParquetWriterSuite + org.apache.spark.sql.comet.CometNativeWriteSuite org.apache.comet.parquet.ParquetReadV1Suite org.apache.comet.parquet.ParquetTimestampLtzAsNtzSuite org.apache.spark.sql.comet.ParquetEncryptionITCase @@ -190,6 +191,7 @@ jobs: org.apache.comet.exec.CometShuffle4_0Suite org.apache.comet.exec.CometNativeColumnarToRowSuite org.apache.comet.exec.CometNativeShuffleSuite + org.apache.comet.exec.CometShuffleInputConversionSuite org.apache.comet.shuffle.CelebornShufflePartitionPusherSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManagerSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornNativeShuffleWriterSuite @@ -207,6 +209,7 @@ jobs: org.apache.comet.exec.CometShuffleManagerSuite org.apache.comet.exec.DisableAQECometShuffleSuite org.apache.spark.shuffle.comet.CometUnboundedShuffleMemoryAllocatorSuite + org.apache.spark.shuffle.comet.CometUnifiedShuffleMemoryAllocatorSuite org.apache.spark.shuffle.sort.SpillSorterSuite - name: "exec" value: | @@ -267,6 +270,7 @@ jobs: org.apache.spark.sql.comet.PlanDataInjectorShuffleLifecycleSuite org.apache.spark.sql.comet.CometDecimalArithmeticViewSuite org.apache.spark.sql.comet.CometDecimalPromotionSuite + org.apache.spark.sql.comet.DecimalPrecisionSuite org.apache.spark.sql.comet.CometScanWithPlanDataSuite org.apache.spark.sql.comet.util.UtilsSuite org.apache.comet.vector.NativeUtilSuite diff --git a/.github/workflows/spark_sql_writer_tests.yml b/.github/workflows/spark_sql_writer_tests.yml index b26ff88e864..6b5fc7f6d65 100644 --- a/.github/workflows/spark_sql_writer_tests.yml +++ b/.github/workflows/spark_sql_writer_tests.yml @@ -115,6 +115,13 @@ jobs: spark-short-version: ${{ inputs.spark-version }} skip-native-build: true + - name: Pre-compile Spark SQL test classes + run: | + cd apache-spark + rm -rf /root/.m2/repository/org/apache/parquet + NOLINT_ON_COMPILE=true build/sbt -Dsbt.log.noformat=true -mem 3072 \ + 'sql/Test/compile' + - name: Run Parquet writer tests run: | cd apache-spark diff --git a/benchmarks/pyspark/run_all_benchmarks.sh b/benchmarks/pyspark/run_all_benchmarks.sh index 452ea8d1ff8..4e43641874a 100755 --- a/benchmarks/pyspark/run_all_benchmarks.sh +++ b/benchmarks/pyspark/run_all_benchmarks.sh @@ -67,6 +67,7 @@ $SPARK_HOME/bin/spark-submit \ # Run Comet JVM shuffle echo "" echo ">>> Running COMET JVM shuffle benchmark..." +# maxBuildSize=-1 measures the hash join alone: no build-side size limit on the rewrite. $SPARK_HOME/bin/spark-submit \ --master "$SPARK_MASTER" \ --executor-memory "$EXECUTOR_MEMORY" \ @@ -86,6 +87,7 @@ $SPARK_HOME/bin/spark-submit \ --conf spark.comet.shuffle.mode=jvm \ --conf spark.comet.shuffle.mode=jvm \ --conf spark.comet.exec.replaceSortMergeJoin=true \ + --conf spark.comet.exec.forceShuffledHashJoin.maxBuildSize=-1 \ --conf spark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager \ --conf spark.sql.extensions=org.apache.comet.CometSparkSessionExtensions \ --conf spark.comet.cast.allowIncompatible=true \ @@ -96,6 +98,7 @@ $SPARK_HOME/bin/spark-submit \ # Run Comet Native shuffle echo "" echo ">>> Running COMET NATIVE shuffle benchmark..." +# maxBuildSize=-1 measures the hash join alone: no build-side size limit on the rewrite. $SPARK_HOME/bin/spark-submit \ --master "$SPARK_MASTER" \ --executor-memory "$EXECUTOR_MEMORY" \ @@ -114,6 +117,7 @@ $SPARK_HOME/bin/spark-submit \ --conf spark.comet.explain.fallback.enabled=true \ --conf spark.comet.shuffle.mode=native \ --conf spark.comet.exec.replaceSortMergeJoin=true \ + --conf spark.comet.exec.forceShuffledHashJoin.maxBuildSize=-1 \ --conf spark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager \ --conf spark.sql.extensions=org.apache.comet.CometSparkSessionExtensions \ --conf spark.comet.cast.allowIncompatible=true \ diff --git a/benchmarks/tpc/engines/comet-hashjoin.toml b/benchmarks/tpc/engines/comet-hashjoin.toml index 202dcad914b..50e184b049b 100644 --- a/benchmarks/tpc/engines/comet-hashjoin.toml +++ b/benchmarks/tpc/engines/comet-hashjoin.toml @@ -31,4 +31,6 @@ driver_class_path = ["$COMET_JAR"] "spark.plugins" = "org.apache.spark.CometPlugin" "spark.shuffle.manager" = "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager" "spark.comet.exec.replaceSortMergeJoin" = "true" +# Measure the hash join alone: no build-side size limit on the rewrite. +"spark.comet.exec.forceShuffledHashJoin.maxBuildSize" = "-1" "spark.comet.expression.Cast.allowIncompatible" = "true" diff --git a/benchmarks/tpc/engines/comet-iceberg-hashjoin.toml b/benchmarks/tpc/engines/comet-iceberg-hashjoin.toml index 421d3201315..4a563613108 100644 --- a/benchmarks/tpc/engines/comet-iceberg-hashjoin.toml +++ b/benchmarks/tpc/engines/comet-iceberg-hashjoin.toml @@ -34,6 +34,8 @@ driver_class_path = ["$COMET_JAR", "$ICEBERG_JAR"] "spark.plugins" = "org.apache.spark.CometPlugin" "spark.shuffle.manager" = "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager" "spark.comet.exec.replaceSortMergeJoin" = "true" +# Measure the hash join alone: no build-side size limit on the rewrite. +"spark.comet.exec.forceShuffledHashJoin.maxBuildSize" = "-1" "spark.comet.expression.Cast.allowIncompatible" = "true" "spark.comet.enabled" = "true" "spark.comet.exec.enabled" = "true" diff --git a/contrib/delta/native/Cargo.lock b/contrib/delta/native/Cargo.lock index c3a4c3289a2..dbbbf9dfc41 100644 --- a/contrib/delta/native/Cargo.lock +++ b/contrib/delta/native/Cargo.lock @@ -488,9 +488,9 @@ dependencies = [ [[package]] name = "datafusion" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "14313430439ef858ed5f52ca9c840e9eea9844da34d3716d05a476449634da5d" +checksum = "4cd5317db77ba6bca25c761a8091b9d294a9f4188e4f9dcd0243e6a57de862ac" dependencies = [ "arrow", "arrow-schema", @@ -534,9 +534,9 @@ dependencies = [ [[package]] name = "datafusion-catalog" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2c0756ebc770bcebae366dfd273dad50935f6103b4d212253cb17e2e53c5427f" +checksum = "f35d36794509b7bfcbf90e11a1b5e4d047123c34610105e2e80a2048351cdb2f" dependencies = [ "arrow", "async-trait", @@ -559,9 +559,9 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bbde2961698705d8402f0f382f08876fc6f1126ffde790ef17833cf0ab87cd59" +checksum = "a6454bfb7f43b83942e2ca2576fd09342eedbd0d323e1eef72c3c8e4d57d7203" dependencies = [ "arrow", "async-trait", @@ -591,9 +591,9 @@ dependencies = [ [[package]] name = "datafusion-common" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a6e06fe711ff815c8cf3143158547aea599702fd43f76f5866c70d5a5b3693a0" +checksum = "d26626e12b224b83351239465dcd00f72d47b02fb5e5b65a643e2dc8ed37f687" dependencies = [ "arrow", "arrow-ipc", @@ -616,9 +616,9 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f7a5107acff4e3ed50a7732acec06bba40e74ab49df021397ed250d60f04e42a" +checksum = "a106420cbfa980f242c22243ae9a6cce661d4d86ce96cff87dacabf7754a4470" dependencies = [ "futures", "log", @@ -627,9 +627,9 @@ dependencies = [ [[package]] name = "datafusion-datasource" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f69f45d2d700785be8df9d4b44410e1c8e4a856f10cc728b418a1066c0b594fb" +checksum = "210aa1e9525e9454f2399dbb8b0b95510145bd8c1fdc35de74e9ccbf81bcb699" dependencies = [ "arrow", "async-trait", @@ -657,9 +657,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "81650af64894ef34b46d2ef818ef244b40a81d90818f3bc642bf8fbf6758f8e0" +checksum = "65b1b041688cad707fa103c3e7f4c6d89b3134ebf67da0129070dad3381851eb" dependencies = [ "arrow", "arrow-ipc", @@ -681,9 +681,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5e2519bed59a6907c79298b9c21fe807bbc34e4ff8414577d699b0741083593" +checksum = "2a7b34e253c4fa5251f174977af910a2f2ec831ec39dbbb56b29ce505cff278b" dependencies = [ "arrow", "async-trait", @@ -704,9 +704,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "635e61290e171597f6e50427e0cb486eda468ab47e880fa24d15bbb0fa24ac7e" +checksum = "1663338f43be6a5746473008e2d0fbde282577ddf0bae90d6c7386e988d71b29" dependencies = [ "arrow", "async-trait", @@ -727,9 +727,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07cd518f85b3a027d913047b4ab0811111a2b734fdf84a51b78234455f10d8fa" +checksum = "7c521ae91b6d899d7190812881f45fe9463c222c11d90149f5332a2e56555904" dependencies = [ "arrow", "arrow-schema", @@ -759,15 +759,15 @@ dependencies = [ [[package]] name = "datafusion-doc" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d44c5919b560862dd1cb20a5c75b3e48c3d9e49719c91d78da5c38eaad5495c4" +checksum = "0c2569c2e74a34fcae3a369ab3912d61e41a062592510f11e31b1dae2725ecdc" [[package]] name = "datafusion-execution" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7620ad535af8730186bdec05d720e40d9a4584b40c7eb75f72c43aa7978e5484" +checksum = "39cfdd5ac0769d1e74a885450e49fde8f63ca7ab41eb479a068564aaea80b837" dependencies = [ "arrow", "arrow-buffer", @@ -791,9 +791,9 @@ dependencies = [ [[package]] name = "datafusion-expr" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ea229af8b73e566fc6ebb1c8dd42ab250f8855c5304bd173f4b50ead78b64304" +checksum = "86560fd7df4dda924f5474b588c56e7ae8e96dcf158c7327c9cfe9daff448ec5" dependencies = [ "arrow", "arrow-schema", @@ -807,14 +807,15 @@ dependencies = [ "datafusion-physical-expr-common", "indexmap", "itertools 0.15.0", + "serde", "serde_json", ] [[package]] name = "datafusion-expr-common" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a18baa840f77760dcc9eaad281329bc57536c5c39e65e8649596f2a51216be7" +checksum = "e82e313cf655cf6bc2dce6cbd23bcb0d20236a54a1b09c436615e4bac4d8c49f" dependencies = [ "arrow", "datafusion-common", @@ -824,9 +825,9 @@ dependencies = [ [[package]] name = "datafusion-functions" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "763bbfeec40b55a8a961d7349c0fcfcec4f3f66cc234d6295250b95371b85e34" +checksum = "939ae2771da8a3fdc23ab4de610ae8fef9515476626b49d96d0d5c782b105771" dependencies = [ "arrow", "arrow-buffer", @@ -852,9 +853,9 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46c0c2b2e856883ecf6a4fa8b23e7c6041183af8a68d0583dc95d1d21e4eaee0" +checksum = "fd6c3aee9f2bc65716787d651e023743a90b4efbc57a2f679f036632e16fb110" dependencies = [ "arrow", "datafusion-common", @@ -873,9 +874,9 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e35e42c87dfbff06934ea0e99cfed3440ac614c3f3bfb8c33be35ca0587d4140" +checksum = "988b5d3dfb4939d81c2f9fab267063610fe52a843620f2fa6cb2a1598e43369c" dependencies = [ "arrow", "datafusion-common", @@ -885,9 +886,9 @@ dependencies = [ [[package]] name = "datafusion-functions-table" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fe2070ecf5638a93c81c8babc7858339ef52165a9f54da1411afd69adf7d85ce" +checksum = "2c767567cce2782d4d0bdfa55fab010aaba0bdbe533509e9e0b9a6a3ee29ed59" dependencies = [ "arrow", "async-trait", @@ -901,9 +902,9 @@ dependencies = [ [[package]] name = "datafusion-functions-window" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0e9a164ebe1998c5c3f9e15c9a0c4e6b599a0c8739b5471bfb089752441963f1" +checksum = "0a311c1f558d1f792f1f631c3670da12ee72a20fbf5d6ca151325d66191c03c1" dependencies = [ "arrow", "datafusion-common", @@ -918,9 +919,9 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "98965b183f92201c16798994820ae032649c4821e089c4c80e7d0f4485d41df5" +checksum = "fffd54b8acf6ecef284cc0677e01e1b8c71be64fec50612ba00a0ff0d9b41a52" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -928,9 +929,9 @@ dependencies = [ [[package]] name = "datafusion-macros" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a711bd5a06ec85d9683e142d28a98ad0679a7220669f5c1b8589e57e9b29e46" +checksum = "14d6df9763c8fb968b2417d797aef4c55ac2e4a895aaef38c72c16a287aa4a40" dependencies = [ "datafusion-doc", "quote", @@ -939,9 +940,9 @@ dependencies = [ [[package]] name = "datafusion-optimizer" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff3256b62ff44616ec83c7f04741d21124970833bb377ed56d9abc5ec6f707ec" +checksum = "acfb8cea6e05e2e75bae183340cf7d53cccb3db0449e2f1fbc07654d980fd19f" dependencies = [ "arrow", "chrono", @@ -958,9 +959,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c1e37a25b5aa1f3a1db114f0ee3605bf01639b83a831dcde35119ee61f0b11d6" +checksum = "9f002cef849e10c0597a217ec7891c2796c0f8b42d739f5d24abc24496373406" dependencies = [ "arrow", "arrow-schema", @@ -980,9 +981,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a4d12787d42e88b26b0d3ec006192fa23ceb832ef2a9e87439ce820a8e038cc" +checksum = "82761522c4536e7059d3b185cd99f10525c2bf48f1bd26a0a57e73e19eadc0fb" dependencies = [ "arrow", "datafusion-common", @@ -995,9 +996,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d13bf7f50fd902f5afc46d50759f5185b5620c40b24054f80f7c648897d3e5a6" +checksum = "f8c5c67aa252886f4d6b331c8114797f5b140464e774c4da36889269a1bf1bd8" dependencies = [ "arrow", "chrono", @@ -1012,9 +1013,9 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11fbe7e6665072957d44ed9587676e17a1f9ab3acfadfc287dac23750585f18a" +checksum = "f1a49dd3a76e4879012682c59dc3735282e60640c97483546a4113564a60bf8f" dependencies = [ "arrow", "datafusion-common", @@ -1031,9 +1032,9 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1265d58e5bce07d154e642a51ff43033b576a6ae40d50a29ac2b9f311004eb52" +checksum = "b1451d6817e5ac0dd21819b19949d1c420ab51d748f55dd9a96ecb5418a24c97" dependencies = [ "arrow", "arrow-data", @@ -1060,15 +1061,16 @@ dependencies = [ "num-traits", "parking_lot", "pin-project-lite", + "serde", "serde_json", "tokio", ] [[package]] name = "datafusion-pruning" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "74b06b333405c05015836ed36cea02cc29d1e2ee8fb414a77532d8346fb56492" +checksum = "f11f0dbe2b5a642f6106e44de8bbd6654fa7220ca14f32313d4efd5514a86968" dependencies = [ "arrow", "datafusion-common", @@ -1082,9 +1084,9 @@ dependencies = [ [[package]] name = "datafusion-session" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e7616b17b1d2fee3e9b13c651eac91b768d153265e6ba88e06e0b4af7883f69c" +checksum = "62dc5ce37a5ce1ff5362949776829aded2529a11dbaa13d24b0ffa6a059e106a" dependencies = [ "arrow-schema", "async-trait", @@ -2119,6 +2121,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" dependencies = [ "serde_core", + "serde_derive", ] [[package]] @@ -2147,7 +2150,6 @@ version = "1.0.151" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" dependencies = [ - "indexmap", "itoa", "memchr", "serde", diff --git a/contrib/delta/native/Cargo.toml b/contrib/delta/native/Cargo.toml index abe73473c4b..43cc1dda273 100644 --- a/contrib/delta/native/Cargo.toml +++ b/contrib/delta/native/Cargo.toml @@ -47,4 +47,4 @@ crate-type = ["rlib"] datafusion-comet-proto = { path = "../../../native/proto" } # For `ExecutionPlan` (return type), `SchemaRef` (the shim hands these in), and `DataFusionError` # (the not-implemented error). Versions/features mirror core so the shared lockfile resolves cleanly. -datafusion = { version = "55.1.0", default-features = false, features = ["parquet"] } +datafusion = { version = "55.2.0", default-features = false, features = ["parquet"] } diff --git a/dev/diffs/3.4.3.diff b/dev/diffs/3.4.3.diff index 5c23919a2ee..af31251b32e 100644 --- a/dev/diffs/3.4.3.diff +++ b/dev/diffs/3.4.3.diff @@ -593,7 +593,7 @@ index f33432ddb6f..b375e285dde 100644 } assert(scanOption.isDefined) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala -index a6b295578d6..1167bbe6554 100644 +index a6b295578d6..1fa898bf840 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala @@ -260,7 +260,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite @@ -606,6 +606,15 @@ index a6b295578d6..1167bbe6554 100644 withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true") { withTempView("df") { val df1 = spark.range(1, 100) +@@ -331,7 +332,7 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite + + val sqlText = "EXPLAIN FORMATTED SELECT (SELECT min(id) FROM df) as v" + val expected_pattern = +- "Subquery:1 Hosting operator id = 2 Hosting Expression = Subquery subquery#x" ++ "Subquery:1 Hosting operator id = \\d+ Hosting Expression = Subquery subquery#x" + + withNormalizedExplain(sqlText) { normalizedOutput => + assert(expected_pattern.r.findAllMatchIn(normalizedOutput).length == 1) @@ -463,7 +464,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite } } diff --git a/dev/diffs/3.5.9.diff b/dev/diffs/3.5.9.diff index d83fe66edce..4ee26ebe93c 100644 --- a/dev/diffs/3.5.9.diff +++ b/dev/diffs/3.5.9.diff @@ -38,7 +38,7 @@ index 49eca0f7555..b9e5d69a6c6 100644 diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala -index daa3ebac992..2b6138c4280 100644 +index daa3ebac992..2c721d85961 100644 --- a/project/SparkBuild.scala +++ b/project/SparkBuild.scala @@ -459,9 +459,11 @@ object SparkBuild extends PomBuild { @@ -53,7 +53,7 @@ index daa3ebac992..2b6138c4280 100644 enable(SparkConnectCommon.settings)(connectCommon) enable(SparkConnect.settings)(connect) -@@ -1198,6 +1200,25 @@ object SQL { +@@ -1198,6 +1200,35 @@ object SQL { ) } @@ -61,16 +61,26 @@ index daa3ebac992..2b6138c4280 100644 + * Comet disables itself when its shuffle manager is not registered. When Comet is enabled, make + * that shuffle manager the default for every test SparkConf in the modules that have Comet on + * their classpath, so that suites which build their own SparkSession or SparkContext run Comet -+ * too, not just the ones that go through SharedSparkSession or TestHive. ++ * too, not just the ones that go through SharedSparkSession or TestHive. Comet's cache serializer ++ * is made the default the same way. Spark keeps the first cache serializer it loads for the rest ++ * of the JVM, so a suite that caches through a session of its own would otherwise leave the ++ * suites after it in Spark's cache format. + */ +object CometTestSettings { + private val cometEnabled = + sys.env.get("ENABLE_COMET").forall(v => v == "1" || v.equalsIgnoreCase("true")) + ++ private val defaults = Seq( ++ "spark.shuffle.manager" -> ++ "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager", ++ "spark.sql.cache.serializer" -> ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") ++ + lazy val settings: Seq[Setting[_]] = -+ if (cometEnabled && !sys.props.contains("spark.shuffle.manager")) { -+ Seq((Test / javaOptions) += -+ "-Dspark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") ++ if (cometEnabled) { ++ defaults.filterNot { case (key, _) => sys.props.contains(key) }.map { ++ case (key, value) => (Test / javaOptions) += s"-D$key=$value" ++ } + } else { + Seq.empty + } @@ -260,10 +270,14 @@ index 0efe0877e9b..423d3b3d76d 100644 -- SELECT_HAVING -- https://github.com/postgres/postgres/blob/REL_12_BETA2/src/test/regress/sql/select_having.sql diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala -index e5494726695..00937f025c2 100644 +index e5494726695..9ceb28855db 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala -@@ -38,7 +38,7 @@ import org.apache.spark.sql.catalyst.util.DateTimeConstants +@@ -35,10 +35,11 @@ import org.apache.spark.sql.catalyst.analysis.TempTableAlreadyExistsException + import org.apache.spark.sql.catalyst.expressions.SubqueryExpression + import org.apache.spark.sql.catalyst.plans.logical.{BROADCAST, Join, JoinStrategyHint, SHUFFLE_HASH} + import org.apache.spark.sql.catalyst.util.DateTimeConstants ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec import org.apache.spark.sql.execution.{ColumnarToRowExec, ExecSubqueryExpression, RDDScanExec, SparkPlan} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEPropagateEmptyRelation} import org.apache.spark.sql.execution.columnar._ @@ -272,7 +286,28 @@ index e5494726695..00937f025c2 100644 import org.apache.spark.sql.execution.ui.SparkListenerSQLAdaptiveExecutionUpdate import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -@@ -519,7 +519,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils +@@ -113,6 +114,10 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + case inMemoryTable @ InMemoryTableScanExec(_, _, relation) => + getNumInMemoryTablesRecursively(relation.cachedPlan) + + getNumInMemoryTablesInSubquery(inMemoryTable) + 1 ++ // Comet's cache scan keeps the predicates pushed into it on its original plan, out of reach ++ // of subquery reuse. The filter above it evaluates the same subqueries, and counts them. ++ case cometScan: CometInMemoryTableScanExec => ++ getNumInMemoryTablesRecursively(cometScan.originalPlan.relation.cachedPlan) + 1 + case p => + getNumInMemoryTablesInSubquery(p) + }.sum +@@ -393,7 +398,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + assert(isExpectStorageLevel(rddId, Disk)) + } + +- test("InMemoryRelation statistics") { ++ test("InMemoryRelation statistics", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + sql("CACHE TABLE testData") + spark.table("testData").queryExecution.withCachedData.collect { + case cached: InMemoryRelation => +@@ -519,7 +525,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils df.collect() } assert( @@ -282,6 +317,18 @@ index e5494726695..00937f025c2 100644 } test("A cached table preserves the partitioning and ordering of its cached SparkPlan") { +@@ -1581,7 +1588,10 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + checkAnswer(spark.table(tableName), Row(LocalDateTime.parse("2021-01-01T00:00:00"))) + spark.table(tableName).queryExecution.withCachedData.collect { + case cached: InMemoryRelation => +- assert(cached.stats.sizeInBytes === 8) ++ // Comet's cache format reports Arrow buffer sizes. ++ if (!isCometEnabled) { ++ assert(cached.stats.sizeInBytes === 8) ++ } + } + sql(s"UNCACHE TABLE $tableName") + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala index 6f3090d8908..52523429c87 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -348,6 +395,30 @@ index 56e9520fdab..917932336df 100644 spark.range(50).write.saveAsTable(s"$dbName.$table1Name") spark.range(100).write.saveAsTable(s"$dbName.$table2Name") +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala +index 747f43fa2a7..1b76251f0bd 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala +@@ -1417,11 +1417,14 @@ class DataFrameSetOperationsSuite extends QueryTest + + val union = df1.union(df2) + union.collect() +- checkIfColumnar(union.queryExecution.executedPlan, +- _.isInstanceOf[InMemoryTableScanExec], supported) +- checkIfColumnar(union.queryExecution.executedPlan, +- _.isInstanceOf[InMemoryTableScanExec], supported) +- checkIfColumnar(union.queryExecution.executedPlan, _.isInstanceOf[UnionExec], supported) ++ // Comet replaces the cache scans and the union with its own operators. ++ if (!isCometEnabled) { ++ checkIfColumnar(union.queryExecution.executedPlan, ++ _.isInstanceOf[InMemoryTableScanExec], supported) ++ checkIfColumnar(union.queryExecution.executedPlan, ++ _.isInstanceOf[InMemoryTableScanExec], supported) ++ checkIfColumnar(union.queryExecution.executedPlan, _.isInstanceOf[UnionExec], supported) ++ } + checkAnswer(union, Row(1) :: Row(2) :: Row(3) :: Row(4) :: Row(5) :: Row(6) :: Nil) + + val nonColumnarUnion = df1.union(Seq(7, 8, 9).toDF("k")) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala index 7ee18df3756..d09f70e5d99 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala @@ -471,6 +542,83 @@ index a1d5d579338..8825683ebcd 100644 case _ => false } case _ => false +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +index bda8c7f2608..10a70745d07 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +@@ -20,6 +20,8 @@ package org.apache.spark.sql + import org.scalatest.concurrent.TimeLimits + import org.scalatest.time.SpanSugar._ + ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec ++import org.apache.spark.sql.execution.SparkPlan + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.columnar.{InMemoryRelation, InMemoryTableScanExec} + import org.apache.spark.sql.functions._ +@@ -35,6 +37,10 @@ class DatasetCacheSuite extends QueryTest + with AdaptiveSparkPlanHelper { + import testImplicits._ + ++ // A scan of a cached relation, Spark's or Comet's. ++ private def isCacheScan(plan: SparkPlan): Boolean = ++ plan.isInstanceOf[InMemoryTableScanExec] || plan.isInstanceOf[CometInMemoryTableScanExec] ++ + /** + * Asserts that a cached [[Dataset]] will be built using the given number of other cached results. + */ +@@ -42,7 +48,7 @@ class DatasetCacheSuite extends QueryTest + val plan = df.queryExecution.withCachedData + assert(plan.isInstanceOf[InMemoryRelation]) + val internalPlan = plan.asInstanceOf[InMemoryRelation].cacheBuilder.cachedPlan +- assert(find(internalPlan)(_.isInstanceOf[InMemoryTableScanExec]).size ++ assert(find(internalPlan)(isCacheScan).size + == numOfCachesDependedUpon) + } + +@@ -252,7 +258,7 @@ class DatasetCacheSuite extends QueryTest + case i: InMemoryRelation => i.cacheBuilder.cachedPlan + } + assert(df2LimitInnerPlan.isDefined && +- !df2LimitInnerPlan.get.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ !df2LimitInnerPlan.get.exists(isCacheScan)) + } + + test("SPARK-27739 Save stats from optimized plan") { +@@ -285,14 +291,14 @@ class DatasetCacheSuite extends QueryTest + val unionDf = df1.union(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.union(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("positive: union by name") { + val unionDf = df1.unionByName(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.unionByName(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by position") { +@@ -300,7 +306,7 @@ class DatasetCacheSuite extends QueryTest + unionDf.cache() + val finalDf = unionDf.union(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by name") { +@@ -308,7 +314,7 @@ class DatasetCacheSuite extends QueryTest + unionDf.cache() + val finalDf = unionDf.unionByName(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala index c4fb4fa943c..a04b23870a8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala @@ -494,10 +642,10 @@ index c4fb4fa943c..a04b23870a8 100644 assert(exchanges.size == 2) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -index f33432ddb6f..b375e285dde 100644 +index f33432ddb6f..4e8522d3aab 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -@@ -22,6 +22,7 @@ import org.scalatest.GivenWhenThen +@@ -22,9 +22,11 @@ import org.scalatest.GivenWhenThen import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression} import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode._ import org.apache.spark.sql.catalyst.plans.ExistenceJoin @@ -505,7 +653,11 @@ index f33432ddb6f..b375e285dde 100644 import org.apache.spark.sql.connector.catalog.{InMemoryTableCatalog, InMemoryTableWithV2FilterCatalog} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive._ -@@ -193,6 +194,7 @@ abstract class DynamicPartitionPruningSuiteBase ++import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec + import org.apache.spark.sql.execution.datasources.v2.BatchScanExec + import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ReusedExchangeExec} + import org.apache.spark.sql.execution.joins.BroadcastHashJoinExec +@@ -193,6 +195,7 @@ abstract class DynamicPartitionPruningSuiteBase } val subqueryBroadcast = dpExprs.collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b @@ -513,7 +665,19 @@ index f33432ddb6f..b375e285dde 100644 } val hasFilter = if (withSubquery) "Should" else "Shouldn't" -@@ -247,6 +249,8 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -229,7 +232,10 @@ abstract class DynamicPartitionPruningSuiteBase + } + + val isMainQueryAdaptive = plan.isInstanceOf[AdaptiveSparkPlanExec] +- subqueriesAll(plan).filterNot(subqueryBroadcast.contains).foreach { s => ++ // Comet's cache scan exposes Spark's own scan as its one subquery, for Spark's UI to draw ++ // the cached plan below it, but the query never runs it as a subquery. ++ subqueriesAll(plan).filterNot(subqueryBroadcast.contains) ++ .filterNot(_.isInstanceOf[InMemoryTableScanExec]).foreach { s => + val subquery = s match { + case r: ReusedSubqueryExec => r.child + case o => o +@@ -247,6 +253,8 @@ abstract class DynamicPartitionPruningSuiteBase val buf = collectDynamicPruningExpressions(df.queryExecution.executedPlan).collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b.index @@ -522,7 +686,7 @@ index f33432ddb6f..b375e285dde 100644 } assert(buf.distinct.size == n) } -@@ -262,6 +266,12 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -262,6 +270,12 @@ abstract class DynamicPartitionPruningSuiteBase case s: BatchScanExec => s.runtimeFilters.collect { case d: DynamicPruningExpression => d.child } @@ -535,7 +699,7 @@ index f33432ddb6f..b375e285dde 100644 case _ => Nil } } -@@ -1204,10 +1214,16 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1204,10 +1218,16 @@ abstract class DynamicPartitionPruningSuiteBase val plan = df.queryExecution.executedPlan val countSubqueryBroadcasts = @@ -554,7 +718,7 @@ index f33432ddb6f..b375e285dde 100644 assert(countSubqueryBroadcasts == 1) assert(countReusedSubqueryBroadcasts == 1) -@@ -1577,6 +1593,7 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1577,6 +1597,7 @@ abstract class DynamicPartitionPruningSuiteBase val subqueryBroadcastExecs = collectWithSubqueries(df.queryExecution.executedPlan) { case s: SubqueryBroadcastExec => s @@ -562,7 +726,7 @@ index f33432ddb6f..b375e285dde 100644 } assert(subqueryBroadcastExecs.size === 1) subqueryBroadcastExecs.foreach { subqueryBroadcastExec => -@@ -1729,6 +1746,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat +@@ -1729,6 +1750,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat case s: BatchScanExec => // we use f1 col for v2 tables due to schema pruning s.output.exists(_.exists(_.argString(maxFields = 100).contains("f1"))) @@ -574,7 +738,7 @@ index f33432ddb6f..b375e285dde 100644 } assert(scanOption.isDefined) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala -index 739557bef30..718fbc5b786 100644 +index 739557bef30..51682036f38 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala @@ -265,7 +265,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite @@ -587,6 +751,15 @@ index 739557bef30..718fbc5b786 100644 withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true") { withTempView("df") { val df1 = spark.range(1, 100) +@@ -336,7 +337,7 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite + + val sqlText = "EXPLAIN FORMATTED SELECT (SELECT min(id) FROM df) as v" + val expected_pattern = +- "Subquery:1 Hosting operator id = 2 Hosting Expression = Subquery subquery#x" ++ "Subquery:1 Hosting operator id = \\d+ Hosting Expression = Subquery subquery#x" + + withNormalizedExplain(sqlText) { normalizedOutput => + assert(expected_pattern.r.findAllMatchIn(normalizedOutput).length == 1) @@ -468,7 +469,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite } } @@ -1045,10 +1218,18 @@ index c26757c9cff..d55775f09d7 100644 protected val baseResourcePath = { // use the same way as `SQLQueryTestSuite` to get the resource path diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala -index 3cf2bfd17ab..5bcf9478e9b 100644 +index 3cf2bfd17ab..11d3ca4d5a9 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala -@@ -1521,7 +1521,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -38,6 +38,7 @@ import org.apache.spark.sql.catalyst.optimizer.{ConvertToLocalRelation, NestedCo + import org.apache.spark.sql.catalyst.parser.ParseException + import org.apache.spark.sql.catalyst.plans.logical.{LocalLimit, Project, RepartitionByExpression, Sort} + import org.apache.spark.sql.connector.catalog.CatalogManager.SESSION_CATALOG_NAME ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.{CommandResultExec, UnionExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.aggregate._ +@@ -1521,7 +1522,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark checkAnswer(sql("select -0.001"), Row(BigDecimal("-0.001"))) } @@ -1058,7 +1239,7 @@ index 3cf2bfd17ab..5bcf9478e9b 100644 AccumulatorSuite.verifyPeakExecutionMemorySet(sparkContext, "external sort") { sql("SELECT * FROM testData2 ORDER BY a ASC, b ASC").collect() } -@@ -1979,8 +1980,15 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -1979,8 +1981,15 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark countAcc.add(1) x }) @@ -1075,7 +1256,15 @@ index 3cf2bfd17ab..5bcf9478e9b 100644 verifyCallCount( df.selectExpr("testUdf(a + 1) + testUdf(1 + a)", "testUdf(a + 1)"), Row(4, 2), 1) -@@ -3750,7 +3758,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -3289,6 +3298,7 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark + + val inMemoryTableScan = collect(queryDf.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c + } + assert(inMemoryTableScan.size == 2) + checkAnswer(queryDf, Row(0, 1) :: Row(1, 2) :: Row(2, 3) :: Row(3, 4) :: Row(4, 5) :: Nil) +@@ -3750,7 +3760,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark } } @@ -1117,7 +1306,7 @@ index 71af1fd69c3..81a04c93c9c 100644 } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala -index 8b4ac474f87..5218e9125c6 100644 +index 8b4ac474f87..b097a14a972 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala @@ -167,7 +167,15 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt @@ -1136,7 +1325,15 @@ index 8b4ac474f87..5218e9125c6 100644 val extensions = create { extensions => // inject rule that will run during AQE query stage preparation and will add custom tags // to the plan -@@ -261,6 +269,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -210,6 +218,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt + df.select("i").filter($"i" > 1).cache() + assert(find(df.filter($"i" > 1).select("i").queryExecution.executedPlan) { + case _: org.apache.spark.sql.execution.columnar.InMemoryTableScanExec => true ++ case _: org.apache.spark.sql.comet.CometInMemoryTableScanExec => true + case _ => false + }.isDefined) + } +@@ -261,6 +270,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } private def testInjectColumnar(enableAQE: Boolean): Unit = { @@ -1144,7 +1341,7 @@ index 8b4ac474f87..5218e9125c6 100644 def collectPlanSteps(plan: SparkPlan): Seq[Int] = plan match { case a: AdaptiveSparkPlanExec => assert(a.toString.startsWith("AdaptiveSparkPlan isFinalPlan=true")) -@@ -316,6 +325,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -316,6 +326,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } test("reset column vectors") { @@ -1152,7 +1349,7 @@ index 8b4ac474f87..5218e9125c6 100644 val session = SparkSession.builder() .master("local[1]") .config(COLUMN_BATCH_SIZE.key, 2) -@@ -484,6 +494,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -484,6 +495,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } test("SPARK-38697: Extend SparkSessionExtensions to inject rules into AQE Optimizer") { @@ -1160,6 +1357,33 @@ index 8b4ac474f87..5218e9125c6 100644 def executedPlan(df: Dataset[java.lang.Long]): SparkPlan = { assert(df.queryExecution.executedPlan.isInstanceOf[AdaptiveSparkPlanExec]) df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec].executedPlan +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +index e827396009d..efbee3975e3 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +@@ -829,14 +829,19 @@ class StatisticsCollectionSuite extends StatisticsCollectionTestBase with Shared + checkTableStats("t1", hasSizeInBytes = true, expectedRowCounts = None) + checkTableStats("t2", hasSizeInBytes = true, expectedRowCounts = None) + assert(getCatalogTable("v1").stats.isEmpty) +- checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ // Comet's cache format reports Arrow buffer sizes for a materialized cached view. ++ if (!isCometEnabled) { ++ checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ } + checkOptimizedPlanStats(spark.table("v2"), 1, None, Seq.empty) + + sql("ANALYZE TABLES COMPUTE STATISTICS") + checkTableStats("t1", hasSizeInBytes = true, expectedRowCounts = Some(10)) + checkTableStats("t2", hasSizeInBytes = true, expectedRowCounts = Some(20)) +- checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) +- checkOptimizedPlanStats(spark.table("v2"), 4, Some(1), Seq.empty) ++ if (!isCometEnabled) { ++ checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ checkOptimizedPlanStats(spark.table("v2"), 4, Some(1), Seq.empty) ++ } + } + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala index 04702201f82..4d38d8d6e51 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala @@ -1397,7 +1621,7 @@ index ae1c0a86a14..1d3b914fd64 100644 // Fail to read ancient datetime values. withSQLConf(SQLConf.PARQUET_REBASE_MODE_IN_READ.key -> EXCEPTION.toString) { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala -index e11191da6a9..005fd4fa657 100644 +index e11191da6a9..0988faf4fcd 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala @@ -51,6 +51,11 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite { @@ -1457,23 +1681,6 @@ index e11191da6a9..005fd4fa657 100644 } } } -@@ -473,10 +484,12 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite { - // aggregate on the other side of the union. - val finalPlan = resultDf.queryExecution.executedPlan - .asInstanceOf[AdaptiveSparkPlanExec].executedPlan -- assert( -- finalPlan.collect { -- case r @ CoalescedShuffleRead() => r -- }.size == 2) -+ if (!isCometEnabled) { -+ assert( -+ finalPlan.collect { -+ case r @ CoalescedShuffleRead() => r -+ }.size == 2) -+ } - } - withSparkSession(test, 100, None) - } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala index 418ca3430bb..eb8267192f8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala @@ -1740,7 +1947,7 @@ index 5a413c77754..207b66e1d7b 100644 import testImplicits._ diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala -index 2f8e401e743..f7923d2b9a7 100644 +index 2f8e401e743..fc4ce35be4b 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala @@ -27,9 +27,11 @@ import org.scalatest.time.SpanSugar._ @@ -2146,7 +2353,18 @@ index 2f8e401e743..f7923d2b9a7 100644 }.size == (if (firstAccess) 2 else 0)) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -2938,7 +2981,8 @@ class AdaptiveQueryExecSuite +@@ -2932,13 +2975,18 @@ class AdaptiveQueryExecSuite + + val (_, adaptive) = runAdaptiveAndVerifyResult( + "SELECT * FROM t1 WHERE c1 < (SELECT c2 FROM t2)") +- assert(findReusedSubquery(adaptive).size == 1) ++ // Comet's cache scan does not plan the subqueries in its pruning predicates, so only the ++ // filter above it runs this one and there is nothing to reuse. ++ if (!isCometEnabled) { ++ assert(findReusedSubquery(adaptive).size == 1) ++ } + } + } } } @@ -2156,7 +2374,7 @@ index 2f8e401e743..f7923d2b9a7 100644 val emptyDf = spark.range(1).where("false") val aggDf1 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df1")) val aggDf2 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df2")) -@@ -2980,7 +3024,9 @@ class AdaptiveQueryExecSuite +@@ -2980,7 +3028,9 @@ class AdaptiveQueryExecSuite val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec]) @@ -2167,6 +2385,147 @@ index 2f8e401e743..f7923d2b9a7 100644 plan.inputPlan.output.zip(plan.finalPhysicalPlan.output).foreach { case (o1, o2) => assert(o1.semanticEquals(o2), "Different output column order after AQE optimization") } +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +index de04938f247..f3351a86946 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +@@ -27,6 +27,7 @@ import org.apache.spark.sql.catalyst.InternalRow + import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSet, In} + import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning + import org.apache.spark.sql.columnar.CachedBatch ++import org.apache.spark.sql.comet.{CometFilterExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{FilterExec, InputAdapter, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions._ +@@ -509,11 +510,17 @@ class InMemoryColumnarQuerySuite extends QueryTest + val planBeforeFilter = collect(df2.queryExecution.executedPlan) { + case f: FilterExec => f.child + case WholeStageCodegenExec(FilterExec(_, i: InputAdapter)) => i.child ++ case f: CometFilterExec => f.child + } +- assert(planBeforeFilter.head.isInstanceOf[InMemoryTableScanExec]) +- + val execPlan = planBeforeFilter.head +- assert(execPlan.executeCollectPublic().length == 0) ++ execPlan match { ++ // Comet's cache scan is columnar only, so count the rows of the batches it emits. ++ case c: CometInMemoryTableScanExec => ++ assert(c.executeColumnar().map(_.numRows().toLong).collect().sum == 0) ++ case _ => ++ assert(execPlan.isInstanceOf[InMemoryTableScanExec]) ++ assert(execPlan.executeCollectPublic().length == 0) ++ } + } + + test("SPARK-25727 - otherCopyArgs in InMemoryRelation does not include outputOrdering") { +@@ -543,7 +550,10 @@ class InMemoryColumnarQuerySuite extends QueryTest + + // InMemoryRelation's stats is updated after materializing RDD + dfFromFile.collect() +- assert(inMemoryRelation.computeStats().sizeInBytes === 16) ++ // Comet's cache format reports Arrow buffer sizes. ++ if (!isCometEnabled) { ++ assert(inMemoryRelation.computeStats().sizeInBytes === 16) ++ } + + // test of catalog table + val dfFromTable = spark.catalog.createTable("table1", workDirPath).cache() +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +index 2c73622739a..5d0efeb263a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +@@ -18,6 +18,7 @@ + package org.apache.spark.sql.execution.columnar + + import org.apache.spark.SparkFunSuite ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.SparkPlan + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions.expr +@@ -37,11 +38,15 @@ class InMemoryRelationSuite extends SparkFunSuite + test("SPARK-47177: Cached SQL plan do not display final AQE plan in explain string") { + def findIMRInnerChild(p: SparkPlan): SparkPlan = { + val tableCache = find(p) { +- case _: InMemoryTableScanExec => true ++ case _: InMemoryTableScanExec | _: CometInMemoryTableScanExec => true + case _ => false + } + assert(tableCache.isDefined) +- tableCache.get.asInstanceOf[InMemoryTableScanExec].relation.innerChildren.head ++ val scan = tableCache.get match { ++ case c: CometInMemoryTableScanExec => c.originalPlan ++ case s => s.asInstanceOf[InMemoryTableScanExec] ++ } ++ scan.relation.innerChildren.head + } + + val d1 = spark.range(1).withColumn("key", expr("id % 100")) +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +index 885286843a1..f0f805a6cd9 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +@@ -17,6 +17,7 @@ + + package org.apache.spark.sql.execution.columnar + ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.internal.SQLConf + import org.apache.spark.sql.test.SharedSparkSession +@@ -181,11 +182,16 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + val result = df.collect().map(_(0)).toArray + assert(result.length === 1) + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- assert(readPartitions === 5) +- assert(readBatches === 10) ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan has none of these test-only accumulators, so there is nothing to count. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ assert(scans.head.readPartitions.value === 5) ++ assert(scans.head.readBatches.value === 10) ++ } + } + + def checkBatchPruning( +@@ -202,14 +208,23 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + df.collect().map(_(0)).toArray + } + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- +- assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution") +- assert( +- readPartitions === expectedReadPartitions, +- s"Wrong number of read partitions: $queryExecution") ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan prunes on the same statistics, but has none of these test-only ++ // accumulators to read, so only the answer above is checked for it. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ val readPartitions = scans.head.readPartitions.value ++ val readBatches = scans.head.readBatches.value ++ assert( ++ readBatches === expectedReadBatches, ++ s"Wrong number of read batches: $queryExecution") ++ assert( ++ readPartitions === expectedReadPartitions, ++ s"Wrong number of read partitions: $queryExecution") ++ } + } + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala index fd52d038ca6..154c800be67 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala @@ -2609,10 +2968,10 @@ index b8f3ea3c6f3..bbd44221288 100644 val workDirPath = workDir.getAbsolutePath val input = spark.range(5).toDF("id") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -index 5479be86e9f..07c81e4b830 100644 +index 5479be86e9f..103dab4657c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -@@ -21,7 +21,7 @@ import scala.reflect.ClassTag +@@ -21,12 +21,13 @@ import scala.reflect.ClassTag import org.apache.spark.AccumulatorSuite import org.apache.spark.internal.config.EXECUTOR_MEMORY @@ -2621,7 +2980,28 @@ index 5479be86e9f..07c81e4b830 100644 import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BitwiseAnd, BitwiseOr, Cast, Expression, Literal, ShiftLeft} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.Inner -@@ -487,7 +487,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + import org.apache.spark.sql.catalyst.plans.logical.BROADCAST + import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, PartitioningCollection} ++import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{DummySparkPlan, SparkPlan, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} + import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec +@@ -157,12 +158,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) +@@ -487,7 +490,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils } } @@ -2631,6 +3011,21 @@ index 5479be86e9f..07c81e4b830 100644 withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") { val t1 = (0 until 100).map(i => (i % 5, i % 13)).toDF("i1", "j1") val t2 = (0 until 100).map(i => (i % 5, i % 14)).toDF("i2", "j2") +@@ -655,12 +659,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala index 5cdbdc27b32..307fba16578 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala @@ -3266,10 +3661,10 @@ index e937173a590..263934fbe7b 100644 spark.internalCreateDataFrame(withoutFilters.execute(), schema) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -index c23bf4204f7..07d215aad2b 100644 +index c23bf4204f7..d922f0b3d23 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -@@ -97,6 +97,19 @@ trait SharedSparkSessionBase +@@ -97,6 +97,24 @@ trait SharedSparkSessionBase // this rule may potentially block testing of other optimization rules such as // ConstantPropagation etc. .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName) @@ -3285,6 +3680,11 @@ index c23bf4204f7..07d215aad2b 100644 + .set("spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") ++ // CometDriverPlugin installs Comet's cache serializer when ++ // spark.comet.exec.inMemoryCache.enabled is on, as it is by default. These sessions do not ++ // load the plugin, so install it here. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } conf.set( StaticSQLConf.WAREHOUSE_PATH, @@ -3370,10 +3770,10 @@ index 6160c3e5f6c..bfc0c618a9b 100644 test("SPARK-4963 DataFrame sample on mutable row return wrong result") { diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala -index 1d646f40b3e..c8192f52f98 100644 +index 1d646f40b3e..b8a043c6248 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala -@@ -53,25 +53,41 @@ object TestHive +@@ -53,25 +53,45 @@ object TestHive new SparkContext( System.getProperty("spark.sql.test.master", "local[1]"), "TestSQLContext", @@ -3425,6 +3825,10 @@ index 1d646f40b3e..c8192f52f98 100644 + .set("spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") ++ ++ // As in SharedSparkSession: what CometDriverPlugin would install. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } + conf diff --git a/dev/diffs/4.0.4.diff b/dev/diffs/4.0.4.diff index f12fdc9e352..895f6bd3db8 100644 --- a/dev/diffs/4.0.4.diff +++ b/dev/diffs/4.0.4.diff @@ -78,7 +78,7 @@ index 046f357ff79..7c6fd199e41 100644 org.apache.datasketches diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala -index 9c2312ed256..b151f6a630c 100644 +index 9c2312ed256..09026e8dc8a 100644 --- a/project/SparkBuild.scala +++ b/project/SparkBuild.scala @@ -414,9 +414,11 @@ object SparkBuild extends PomBuild { @@ -93,7 +93,7 @@ index 9c2312ed256..b151f6a630c 100644 enable(HiveThriftServer.settings)(hiveThriftServer) -@@ -1183,6 +1185,25 @@ object SQL { +@@ -1183,6 +1185,35 @@ object SQL { } } @@ -101,16 +101,26 @@ index 9c2312ed256..b151f6a630c 100644 + * Comet disables itself when its shuffle manager is not registered. When Comet is enabled, make + * that shuffle manager the default for every test SparkConf in the modules that have Comet on + * their classpath, so that suites which build their own SparkSession or SparkContext run Comet -+ * too, not just the ones that go through SharedSparkSession or TestHive. ++ * too, not just the ones that go through SharedSparkSession or TestHive. Comet's cache serializer ++ * is made the default the same way. Spark keeps the first cache serializer it loads for the rest ++ * of the JVM, so a suite that caches through a session of its own would otherwise leave the ++ * suites after it in Spark's cache format. + */ +object CometTestSettings { + private val cometEnabled = + sys.env.get("ENABLE_COMET").forall(v => v == "1" || v.equalsIgnoreCase("true")) + ++ private val defaults = Seq( ++ "spark.shuffle.manager" -> ++ "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager", ++ "spark.sql.cache.serializer" -> ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") ++ + lazy val settings: Seq[Setting[_]] = -+ if (cometEnabled && !sys.props.contains("spark.shuffle.manager")) { -+ Seq((Test / javaOptions) += -+ "-Dspark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") ++ if (cometEnabled) { ++ defaults.filterNot { case (key, _) => sys.props.contains(key) }.map { ++ case (key, value) => (Test / javaOptions) += s"-D$key=$value" ++ } + } else { + Seq.empty + } @@ -375,10 +385,14 @@ index 21a3ce1e122..f4762ab98f0 100644 -- In COMPENSATION views get invalidated if the type can't cast diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala -index 0f42502f1d9..e9ff802141f 100644 +index 0f42502f1d9..98479690165 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala -@@ -39,7 +39,7 @@ import org.apache.spark.sql.catalyst.util.DateTimeConstants +@@ -36,10 +36,11 @@ import org.apache.spark.sql.catalyst.analysis.TempTableAlreadyExistsException + import org.apache.spark.sql.catalyst.expressions.SubqueryExpression + import org.apache.spark.sql.catalyst.plans.logical.{BROADCAST, Join, JoinStrategyHint, SHUFFLE_HASH} + import org.apache.spark.sql.catalyst.util.DateTimeConstants ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec import org.apache.spark.sql.execution.{ColumnarToRowExec, ExecSubqueryExpression, RDDScanExec, SparkPlan, SparkPlanInfo} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEPropagateEmptyRelation} import org.apache.spark.sql.execution.columnar._ @@ -387,7 +401,28 @@ index 0f42502f1d9..e9ff802141f 100644 import org.apache.spark.sql.execution.ui.SparkListenerSQLAdaptiveExecutionUpdate import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -@@ -520,7 +520,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils +@@ -114,6 +115,10 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + case inMemoryTable @ InMemoryTableScanExec(_, _, relation) => + getNumInMemoryTablesRecursively(relation.cachedPlan) + + getNumInMemoryTablesInSubquery(inMemoryTable) + 1 ++ // Comet's cache scan keeps the predicates pushed into it on its original plan, out of reach ++ // of subquery reuse. The filter above it evaluates the same subqueries, and counts them. ++ case cometScan: CometInMemoryTableScanExec => ++ getNumInMemoryTablesRecursively(cometScan.originalPlan.relation.cachedPlan) + 1 + case p => + getNumInMemoryTablesInSubquery(p) + }.sum +@@ -394,7 +399,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + assert(isExpectStorageLevel(rddId, Disk)) + } + +- test("InMemoryRelation statistics") { ++ test("InMemoryRelation statistics", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + sql("CACHE TABLE testData") + spark.table("testData").queryExecution.withCachedData.collect { + case cached: InMemoryRelation => +@@ -520,7 +526,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils df.collect() } assert( @@ -397,7 +432,19 @@ index 0f42502f1d9..e9ff802141f 100644 } test("A cached table preserves the partitioning and ordering of its cached SparkPlan") { -@@ -1659,9 +1660,18 @@ class CachedTableSuite extends QueryTest with SQLTestUtils +@@ -1588,7 +1595,10 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + checkAnswer(spark.table(tableName), Row(LocalDateTime.parse("2021-01-01T00:00:00"))) + spark.table(tableName).queryExecution.withCachedData.collect { + case cached: InMemoryRelation => +- assert(cached.stats.sizeInBytes === 8) ++ // Comet's cache format reports Arrow buffer sizes. ++ if (!isCometEnabled) { ++ assert(cached.stats.sizeInBytes === 8) ++ } + } + sql(s"UNCACHE TABLE $tableName") + } +@@ -1659,9 +1669,19 @@ class CachedTableSuite extends QueryTest with SQLTestUtils _.nodeName.contains("TableCacheQueryStage")) val aqeNode = findNodeInSparkPlanInfo(inMemoryScanNode.get, _.nodeName.contains("AdaptiveSparkPlan")) @@ -405,20 +452,29 @@ index 0f42502f1d9..e9ff802141f 100644 - _.nodeName.contains("ResultQueryStage")) - aqePlanRoot.get.children.head.nodeName == "AQEShuffleRead" + // Spark 4.0 wraps results in ResultQueryStage. The coalescing indicator is AQEShuffleRead -+ // as the direct child of InputAdapter. ++ // as its direct child, as in Spark, when the table is cached in Comet's columnar format, ++ // or as the direct child of InputAdapter when Comet's plan is cached in Spark's format. + // AdaptiveSparkPlan -> ResultQueryStage -> WholestageCodegen -> + // CometColumnarToRow -> InputAdapter -> AQEShuffleRead (if coalesced) + val resultStage = aqeNode.get.children.head // ResultQueryStage -+ val wsc = resultStage.children.head // WholeStageCodegen ++ val wsc = resultStage.children.head // WholeStageCodegen or AQEShuffleRead + val c2r = wsc.children.head // ColumnarToRow or CometColumnarToRow + val inputAdapter = c2r.children.head // InputAdapter -+ resultStage.nodeName == "ResultQueryStage" && ++ resultStage.nodeName == "ResultQueryStage" && (wsc.nodeName == "AQEShuffleRead" || + wsc.nodeName.startsWith("WholeStageCodegen") && // could be "WholeStageCodegen (1)" + (c2r.nodeName == "CometColumnarToRow" || c2r.nodeName == "ColumnarToRow") && -+ inputAdapter.children.head.nodeName == "AQEShuffleRead" ++ inputAdapter.children.head.nodeName == "AQEShuffleRead") } withTempView("t0", "t1", "t2") { +@@ -1750,6 +1770,7 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + val cached = spark.table("t") + val tableCache = collect(cached.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c.originalPlan + } + if (expected == StorageLevel.NONE) { + assert(tableCache.isEmpty) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala index 9db406ff12f..e59a931c06a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -485,6 +541,30 @@ index ed182322aec..1ae6afa686a 100644 spark.range(50).write.saveAsTable(s"$dbName.$table1Name") spark.range(100).write.saveAsTable(s"$dbName.$table2Name") +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala +index 332be4c7bbc..ec66595e873 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala +@@ -1493,11 +1493,14 @@ class DataFrameSetOperationsSuite extends QueryTest + + val union = df1.union(df2) + union.collect() +- checkIfColumnar(union.queryExecution.executedPlan, +- _.isInstanceOf[InMemoryTableScanExec], supported) +- checkIfColumnar(union.queryExecution.executedPlan, +- _.isInstanceOf[InMemoryTableScanExec], supported) +- checkIfColumnar(union.queryExecution.executedPlan, _.isInstanceOf[UnionExec], supported) ++ // Comet replaces the cache scans and the union with its own operators. ++ if (!isCometEnabled) { ++ checkIfColumnar(union.queryExecution.executedPlan, ++ _.isInstanceOf[InMemoryTableScanExec], supported) ++ checkIfColumnar(union.queryExecution.executedPlan, ++ _.isInstanceOf[InMemoryTableScanExec], supported) ++ checkIfColumnar(union.queryExecution.executedPlan, _.isInstanceOf[UnionExec], supported) ++ } + checkAnswer(union, Row(1) :: Row(2) :: Row(3) :: Row(4) :: Row(5) :: Row(6) :: Nil) + + val nonColumnarUnion = df1.union(Seq(7, 8, 9).toDF("k")) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala index d9ce3000a0c..f2d044ed6b8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala @@ -608,6 +688,97 @@ index 552e2b2e274..17a5ae20f0f 100644 case _ => false } case _ => false +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +index 9d8aaf8d90e..41afdbdafc1 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +@@ -20,6 +20,8 @@ package org.apache.spark.sql + import org.scalatest.concurrent.TimeLimits + import org.scalatest.time.SpanSugar._ + ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec ++import org.apache.spark.sql.execution.SparkPlan + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.columnar.{InMemoryRelation, InMemoryTableScanExec} + import org.apache.spark.sql.functions._ +@@ -36,6 +38,10 @@ class DatasetCacheSuite extends QueryTest + with AdaptiveSparkPlanHelper { + import testImplicits._ + ++ // A scan of a cached relation, Spark's or Comet's. ++ private def isCacheScan(plan: SparkPlan): Boolean = ++ plan.isInstanceOf[InMemoryTableScanExec] || plan.isInstanceOf[CometInMemoryTableScanExec] ++ + /** + * Asserts that a cached [[Dataset]] will be built using the given number of other cached results. + */ +@@ -43,7 +49,7 @@ class DatasetCacheSuite extends QueryTest + val plan = df.queryExecution.withCachedData + assert(plan.isInstanceOf[InMemoryRelation]) + val internalPlan = plan.asInstanceOf[InMemoryRelation].cacheBuilder.cachedPlan +- assert(find(internalPlan)(_.isInstanceOf[InMemoryTableScanExec]).size ++ assert(find(internalPlan)(isCacheScan).size + == numOfCachesDependedUpon) + } + +@@ -253,7 +259,7 @@ class DatasetCacheSuite extends QueryTest + case i: InMemoryRelation => i.cacheBuilder.cachedPlan + } + assert(df2LimitInnerPlan.isDefined && +- !df2LimitInnerPlan.get.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ !df2LimitInnerPlan.get.exists(isCacheScan)) + } + + test("SPARK-27739 Save stats from optimized plan") { +@@ -286,14 +292,14 @@ class DatasetCacheSuite extends QueryTest + val unionDf = df1.union(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.union(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("positive: union by name") { + val unionDf = df1.unionByName(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.unionByName(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by position") { +@@ -301,7 +307,7 @@ class DatasetCacheSuite extends QueryTest + unionDf.cache() + val finalDf = unionDf.union(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by name") { +@@ -309,7 +315,7 @@ class DatasetCacheSuite extends QueryTest + unionDf.cache() + val finalDf = unionDf.unionByName(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + } + } +@@ -321,11 +327,11 @@ class DatasetCacheSuite extends QueryTest + df1.cache() + // This is exactly the same as df1. + val df2 = spark.range(5).select(struct($"id".as("name", metadata))) +- assert(df2.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(df2.queryExecution.executedPlan.exists(isCacheScan)) + + val metadata2 = Metadata.fromJson("""{"k2": "v2"}""") + // Same with df1 except for the Alias metadata + val df3 = spark.range(5).select(struct($"id".as("name", metadata2))) +- assert(!df3.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!df3.queryExecution.executedPlan.exists(isCacheScan)) + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala index 81713c777bc..b5f92ed9742 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala @@ -631,10 +802,10 @@ index 81713c777bc..b5f92ed9742 100644 assert(exchanges.size == 2) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -index 2c24cc7d570..12d897866da 100644 +index 2c24cc7d570..505157acd91 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -@@ -22,6 +22,7 @@ import org.scalatest.GivenWhenThen +@@ -22,9 +22,11 @@ import org.scalatest.GivenWhenThen import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression} import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode._ import org.apache.spark.sql.catalyst.plans.ExistenceJoin @@ -642,7 +813,11 @@ index 2c24cc7d570..12d897866da 100644 import org.apache.spark.sql.connector.catalog.{InMemoryTableCatalog, InMemoryTableWithV2FilterCatalog} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive._ -@@ -193,6 +194,7 @@ abstract class DynamicPartitionPruningSuiteBase ++import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec + import org.apache.spark.sql.execution.datasources.v2.BatchScanExec + import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ReusedExchangeExec} + import org.apache.spark.sql.execution.joins.BroadcastHashJoinExec +@@ -193,6 +195,7 @@ abstract class DynamicPartitionPruningSuiteBase } val subqueryBroadcast = dpExprs.collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b @@ -650,7 +825,19 @@ index 2c24cc7d570..12d897866da 100644 } val hasFilter = if (withSubquery) "Should" else "Shouldn't" -@@ -247,6 +249,8 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -229,7 +232,10 @@ abstract class DynamicPartitionPruningSuiteBase + } + + val isMainQueryAdaptive = plan.isInstanceOf[AdaptiveSparkPlanExec] +- subqueriesAll(plan).filterNot(subqueryBroadcast.contains).foreach { s => ++ // Comet's cache scan exposes Spark's own scan as its one subquery, for Spark's UI to draw ++ // the cached plan below it, but the query never runs it as a subquery. ++ subqueriesAll(plan).filterNot(subqueryBroadcast.contains) ++ .filterNot(_.isInstanceOf[InMemoryTableScanExec]).foreach { s => + val subquery = s match { + case r: ReusedSubqueryExec => r.child + case o => o +@@ -247,6 +253,8 @@ abstract class DynamicPartitionPruningSuiteBase val buf = collectDynamicPruningExpressions(df.queryExecution.executedPlan).collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b.indices.map(idx => b.buildKeys(idx)) @@ -659,7 +846,7 @@ index 2c24cc7d570..12d897866da 100644 } assert(buf.distinct.size == n) } -@@ -262,6 +266,12 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -262,6 +270,12 @@ abstract class DynamicPartitionPruningSuiteBase case s: BatchScanExec => s.runtimeFilters.collect { case d: DynamicPruningExpression => d.child } @@ -672,7 +859,7 @@ index 2c24cc7d570..12d897866da 100644 case _ => Nil } } -@@ -1204,10 +1214,16 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1204,10 +1218,16 @@ abstract class DynamicPartitionPruningSuiteBase val plan = df.queryExecution.executedPlan val countSubqueryBroadcasts = @@ -691,7 +878,7 @@ index 2c24cc7d570..12d897866da 100644 assert(countSubqueryBroadcasts == 1) assert(countReusedSubqueryBroadcasts == 1) -@@ -1578,6 +1594,7 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1578,6 +1598,7 @@ abstract class DynamicPartitionPruningSuiteBase val subqueryBroadcastExecs = collectWithSubqueries(df.queryExecution.executedPlan) { case s: SubqueryBroadcastExec => s @@ -699,7 +886,7 @@ index 2c24cc7d570..12d897866da 100644 } assert(subqueryBroadcastExecs.size === 1) subqueryBroadcastExecs.foreach { subqueryBroadcastExec => -@@ -1730,6 +1747,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat +@@ -1730,6 +1751,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat case s: BatchScanExec => // we use f1 col for v2 tables due to schema pruning s.output.exists(_.exists(_.argString(maxFields = 100).contains("f1"))) @@ -711,7 +898,7 @@ index 2c24cc7d570..12d897866da 100644 } assert(scanOption.isDefined) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala -index ef9c301a66a..49fad8cbbf0 100644 +index ef9c301a66a..d97e9f76dc3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala @@ -268,7 +268,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite @@ -724,6 +911,15 @@ index ef9c301a66a..49fad8cbbf0 100644 withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true") { withTempView("df") { val df1 = spark.range(1, 100) +@@ -339,7 +340,7 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite + + val sqlText = "EXPLAIN FORMATTED SELECT (SELECT min(id) FROM df) as v" + val expected_pattern = +- "Subquery:1 Hosting operator id = 2 Hosting Expression = Subquery subquery#x" ++ "Subquery:1 Hosting operator id = \\d+ Hosting Expression = Subquery subquery#x" + + withNormalizedExplain(sqlText) { normalizedOutput => + assert(expected_pattern.r.findAllMatchIn(normalizedOutput).length == 1) @@ -471,7 +472,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite } } @@ -1192,10 +1388,18 @@ index ad424b3a7cc..4ece0117a34 100644 protected val baseResourcePath = { // use the same way as `SQLQueryTestSuite` to get the resource path diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala -index f294ff81021..02d72be8d29 100644 +index f294ff81021..f019c6de8e9 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala -@@ -1524,7 +1524,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -38,6 +38,7 @@ import org.apache.spark.sql.catalyst.optimizer.{ConvertToLocalRelation, NestedCo + import org.apache.spark.sql.catalyst.parser.ParseException + import org.apache.spark.sql.catalyst.plans.logical.{LocalLimit, Project, RepartitionByExpression, Sort} + import org.apache.spark.sql.connector.catalog.CatalogManager.SESSION_CATALOG_NAME ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.{CommandResultExec, UnionExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.aggregate._ +@@ -1524,7 +1525,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark checkAnswer(sql("select -0.001"), Row(BigDecimal("-0.001"))) } @@ -1205,7 +1409,7 @@ index f294ff81021..02d72be8d29 100644 AccumulatorSuite.verifyPeakExecutionMemorySet(sparkContext, "external sort") { sql("SELECT * FROM testData2 ORDER BY a ASC, b ASC").collect() } -@@ -1985,8 +1986,15 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -1985,8 +1987,15 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark countAcc.add(1) x }) @@ -1222,6 +1426,24 @@ index f294ff81021..02d72be8d29 100644 verifyCallCount( df.selectExpr("testUdf(a + 1) + testUdf(1 + a)", "testUdf(a + 1)"), Row(4, 2), 1) +@@ -3278,6 +3287,7 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark + + val inMemoryTableScan = collect(queryDf.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c + } + assert(inMemoryTableScan.size == 2) + checkAnswer(queryDf, Row(0, 1) :: Row(1, 2) :: Row(2, 3) :: Row(3, 4) :: Row(4, 5) :: Nil) +@@ -4744,7 +4754,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark + ) + } + +- test("SPARK-47939: Explain should work with parameterized queries") { ++ test("SPARK-47939: Explain should work with parameterized queries", ++ IgnoreComet("Comet explain output is different")) { + def checkQueryPlan(df: DataFrame, plan: String): Unit = assert( + df.collect() + .map(_.getString(0)) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala index 575a4ae69d1..129d9f27232 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala @@ -1254,7 +1476,7 @@ index 575a4ae69d1..129d9f27232 100644 } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala -index c1c041509c3..d068241d585 100644 +index c1c041509c3..6fc92c40ae7 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala @@ -179,7 +179,15 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt @@ -1273,7 +1495,15 @@ index c1c041509c3..d068241d585 100644 val extensions = create { extensions => // inject rule that will run during AQE query stage preparation and will add custom tags // to the plan -@@ -273,6 +281,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -222,6 +230,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt + df.select("i").filter($"i" > 1).cache() + assert(find(df.filter($"i" > 1).select("i").queryExecution.executedPlan) { + case _: org.apache.spark.sql.execution.columnar.InMemoryTableScanExec => true ++ case _: org.apache.spark.sql.comet.CometInMemoryTableScanExec => true + case _ => false + }.isDefined) + } +@@ -273,6 +282,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } private def testInjectColumnar(enableAQE: Boolean): Unit = { @@ -1281,7 +1511,7 @@ index c1c041509c3..d068241d585 100644 def collectPlanSteps(plan: SparkPlan): Seq[Int] = plan match { case a: AdaptiveSparkPlanExec => assert(a.toString.startsWith("AdaptiveSparkPlan isFinalPlan=true")) -@@ -328,6 +337,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -328,6 +338,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } test("reset column vectors") { @@ -1289,7 +1519,7 @@ index c1c041509c3..d068241d585 100644 val session = SparkSession.builder() .master("local[1]") .config(COLUMN_BATCH_SIZE.key, 2) -@@ -496,6 +506,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -496,6 +507,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } test("SPARK-38697: Extend SparkSessionExtensions to inject rules into AQE Optimizer") { @@ -1312,6 +1542,33 @@ index 5ba69c8f9d9..ac1256afe88 100644 sc = new SparkContext("local[2]", "test") val session = classic.SparkSession.builder().sparkContext(sc).getOrCreate() import session.implicits._ +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +index 5222d5ce266..6f8f8bec3d4 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +@@ -847,14 +847,19 @@ class StatisticsCollectionSuite extends StatisticsCollectionTestBase with Shared + checkTableStats("t1", hasSizeInBytes = true, expectedRowCounts = None) + checkTableStats("t2", hasSizeInBytes = true, expectedRowCounts = None) + assert(getCatalogTable("v1").stats.isEmpty) +- checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ // Comet's cache format reports Arrow buffer sizes for a materialized cached view. ++ if (!isCometEnabled) { ++ checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ } + checkOptimizedPlanStats(spark.table("v2"), 1, None, Seq.empty) + + sql("ANALYZE TABLES COMPUTE STATISTICS") + checkTableStats("t1", hasSizeInBytes = true, expectedRowCounts = Some(10)) + checkTableStats("t2", hasSizeInBytes = true, expectedRowCounts = Some(20)) +- checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) +- checkOptimizedPlanStats(spark.table("v2"), 4, Some(1), Seq.empty) ++ if (!isCometEnabled) { ++ checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ checkOptimizedPlanStats(spark.table("v2"), 4, Some(1), Seq.empty) ++ } + } + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala index 0df7f806272..9cdfe8b8f46 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala @@ -1511,7 +1768,7 @@ index a40e34d94d0..abc1f035d15 100644 } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala -index 11e9547dfc5..ba340c4ebcf 100644 +index 11e9547dfc5..327e0b1bc5b 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala @@ -24,6 +24,8 @@ import org.apache.spark.sql.{AnalysisException, Row} @@ -1519,7 +1776,7 @@ index 11e9547dfc5..ba340c4ebcf 100644 import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.util.CollationFactory +import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashJoinExec, CometSortMergeJoinExec} -+import org.apache.spark.sql.comet.CometHashAggregateExec ++import org.apache.spark.sql.comet.{CometHashAggregateExec, CometInMemoryTableScanExec} import org.apache.spark.sql.connector.{DatasourceV2SQLBase, FakeV2ProviderWithCustomSchema} import org.apache.spark.sql.connector.catalog.{Identifier, InMemoryTable} import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.CatalogHelper @@ -1578,6 +1835,15 @@ index 11e9547dfc5..ba340c4ebcf 100644 }.head.isInstanceOf[ArrayTransform]) } } +@@ -1897,7 +1909,7 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper { + // Checks in-memory fetching code path. + val all = sql("SELECT col FROM tbl") + assert(all.queryExecution.executedPlan.collectFirst { +- case _: InMemoryTableScanExec => true ++ case _: InMemoryTableScanExec | _: CometInMemoryTableScanExec => true + }.nonEmpty) + checkAnswer(all, Row("a")) + // Checks column stats code path. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala index 3eeed2e4175..9f21d547c1c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala @@ -1760,7 +2026,7 @@ index 04d33ecd3d5..450df347297 100644 // Fail to read ancient datetime values. withSQLConf(SQLConf.PARQUET_REBASE_MODE_IN_READ.key -> EXCEPTION.toString) { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala -index da43b0cfc58..ca1f6f47a9f 100644 +index da43b0cfc58..653cd7036d1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala @@ -54,6 +54,11 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper @@ -1820,23 +2086,6 @@ index da43b0cfc58..ca1f6f47a9f 100644 } } } -@@ -468,10 +479,12 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper - // Shuffle partition coalescing of the join is performed independent of the non-grouping - // aggregate on the other side of the union. - val finalPlan = stripAQEPlan(resultDf.queryExecution.executedPlan) -- assert( -- finalPlan.collect { -- case r @ CoalescedShuffleRead() => r -- }.size == 2) -+ if (!isCometEnabled) { -+ assert( -+ finalPlan.collect { -+ case r @ CoalescedShuffleRead() => r -+ }.size == 2) -+ } - } - withSparkSession(test, 100, None) - } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala index 418ca3430bb..eb8267192f8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala @@ -2387,7 +2636,7 @@ index a3cfdc5a240..3793b6191bf 100644 }) checkAnswer(distinctWithId, Seq(Row(1, 0), Row(1, 0))) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala -index fb8fab6a80f..c5768935805 100644 +index fb8fab6a80f..cfcff43d557 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala @@ -28,12 +28,14 @@ import org.apache.spark.SparkException @@ -2872,7 +3121,18 @@ index fb8fab6a80f..c5768935805 100644 }.isEmpty) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -3048,7 +3101,8 @@ class AdaptiveQueryExecSuite +@@ -3042,13 +3095,18 @@ class AdaptiveQueryExecSuite + + val (_, adaptive) = runAdaptiveAndVerifyResult( + "SELECT * FROM t1 WHERE c1 < (SELECT c2 FROM t2)") +- assert(findReusedSubquery(adaptive).size == 1) ++ // Comet's cache scan does not plan the subqueries in its pruning predicates, so only the ++ // filter above it runs this one and there is nothing to reuse. ++ if (!isCometEnabled) { ++ assert(findReusedSubquery(adaptive).size == 1) ++ } + } + } } } @@ -2882,7 +3142,7 @@ index fb8fab6a80f..c5768935805 100644 val emptyDf = spark.range(1).where("false") val aggDf1 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df1")) val aggDf2 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df2")) -@@ -3138,7 +3192,8 @@ class AdaptiveQueryExecSuite +@@ -3138,7 +3196,8 @@ class AdaptiveQueryExecSuite val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec]) @@ -2892,6 +3152,147 @@ index fb8fab6a80f..c5768935805 100644 plan.inputPlan.output.zip(plan.finalPhysicalPlan.output).foreach { case (o1, o2) => assert(o1.semanticEquals(o2), "Different output column order after AQE optimization") } +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +index 4f07d3d1c03..9610c703610 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +@@ -28,6 +28,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, + import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning + import org.apache.spark.sql.classic.DataFrame + import org.apache.spark.sql.columnar.CachedBatch ++import org.apache.spark.sql.comet.{CometFilterExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{FilterExec, InputAdapter, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions._ +@@ -510,11 +511,17 @@ class InMemoryColumnarQuerySuite extends QueryTest + val planBeforeFilter = collect(df2.queryExecution.executedPlan) { + case f: FilterExec => f.child + case WholeStageCodegenExec(FilterExec(_, i: InputAdapter)) => i.child ++ case f: CometFilterExec => f.child + } +- assert(planBeforeFilter.head.isInstanceOf[InMemoryTableScanExec]) +- + val execPlan = planBeforeFilter.head +- assert(execPlan.executeCollectPublic().length == 0) ++ execPlan match { ++ // Comet's cache scan is columnar only, so count the rows of the batches it emits. ++ case c: CometInMemoryTableScanExec => ++ assert(c.executeColumnar().map(_.numRows().toLong).collect().sum == 0) ++ case _ => ++ assert(execPlan.isInstanceOf[InMemoryTableScanExec]) ++ assert(execPlan.executeCollectPublic().length == 0) ++ } + } + + test("SPARK-25727 - otherCopyArgs in InMemoryRelation does not include outputOrdering") { +@@ -544,7 +551,10 @@ class InMemoryColumnarQuerySuite extends QueryTest + + // InMemoryRelation's stats is updated after materializing RDD + dfFromFile.collect() +- assert(inMemoryRelation.computeStats().sizeInBytes === 16) ++ // Comet's cache format reports Arrow buffer sizes. ++ if (!isCometEnabled) { ++ assert(inMemoryRelation.computeStats().sizeInBytes === 16) ++ } + + // test of catalog table + val dfFromTable = spark.catalog.createTable("table1", workDirPath).cache() +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +index 2c73622739a..5d0efeb263a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +@@ -18,6 +18,7 @@ + package org.apache.spark.sql.execution.columnar + + import org.apache.spark.SparkFunSuite ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.SparkPlan + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions.expr +@@ -37,11 +38,15 @@ class InMemoryRelationSuite extends SparkFunSuite + test("SPARK-47177: Cached SQL plan do not display final AQE plan in explain string") { + def findIMRInnerChild(p: SparkPlan): SparkPlan = { + val tableCache = find(p) { +- case _: InMemoryTableScanExec => true ++ case _: InMemoryTableScanExec | _: CometInMemoryTableScanExec => true + case _ => false + } + assert(tableCache.isDefined) +- tableCache.get.asInstanceOf[InMemoryTableScanExec].relation.innerChildren.head ++ val scan = tableCache.get match { ++ case c: CometInMemoryTableScanExec => c.originalPlan ++ case s => s.asInstanceOf[InMemoryTableScanExec] ++ } ++ scan.relation.innerChildren.head + } + + val d1 = spark.range(1).withColumn("key", expr("id % 100")) +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +index 88ff51d0ff4..7d00fca3b0a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +@@ -17,6 +17,7 @@ + + package org.apache.spark.sql.execution.columnar + ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.internal.SQLConf + import org.apache.spark.sql.test.SharedSparkSession +@@ -181,11 +182,16 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + val result = df.collect().map(_(0)).toArray + assert(result.length === 1) + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- assert(readPartitions === 5) +- assert(readBatches === 10) ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan has none of these test-only accumulators, so there is nothing to count. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ assert(scans.head.readPartitions.value === 5) ++ assert(scans.head.readBatches.value === 10) ++ } + } + + def checkBatchPruning( +@@ -202,14 +208,23 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + df.collect().map(_(0)).toArray + } + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- +- assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution") +- assert( +- readPartitions === expectedReadPartitions, +- s"Wrong number of read partitions: $queryExecution") ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan prunes on the same statistics, but has none of these test-only ++ // accumulators to read, so only the answer above is checked for it. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ val readPartitions = scans.head.readPartitions.value ++ val readBatches = scans.head.readBatches.value ++ assert( ++ readBatches === expectedReadBatches, ++ s"Wrong number of read batches: $queryExecution") ++ assert( ++ readPartitions === expectedReadPartitions, ++ s"Wrong number of read partitions: $queryExecution") ++ } + } + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala index 269990d7d14..140ee4112b1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala @@ -3334,10 +3735,10 @@ index b8f3ea3c6f3..bbd44221288 100644 val workDirPath = workDir.getAbsolutePath val input = spark.range(5).toDF("id") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -index 69dd04e07d5..781018ecc66 100644 +index 69dd04e07d5..a78dd593985 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -@@ -21,7 +21,7 @@ import scala.reflect.ClassTag +@@ -21,13 +21,14 @@ import scala.reflect.ClassTag import org.apache.spark.AccumulatorSuite import org.apache.spark.internal.config.EXECUTOR_MEMORY @@ -3346,7 +3747,29 @@ index 69dd04e07d5..781018ecc66 100644 import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BitwiseAnd, BitwiseOr, Cast, Expression, Literal, ShiftLeft} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.Inner -@@ -488,7 +488,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + import org.apache.spark.sql.catalyst.plans.logical.BROADCAST + import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, PartitioningCollection} + import org.apache.spark.sql.classic.{Dataset, SparkSession} ++import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{DummySparkPlan, SparkPlan, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} + import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec +@@ -158,12 +159,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) +@@ -488,7 +491,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils } } @@ -3356,6 +3779,21 @@ index 69dd04e07d5..781018ecc66 100644 withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") { val t1 = (0 until 100).map(i => (i % 5, i % 13)).toDF("i1", "j1") val t2 = (0 until 100).map(i => (i % 5, i % 14)).toDF("i2", "j2") +@@ -656,12 +660,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala index 0dd90925d3c..7d53ec845ef 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala @@ -4058,10 +4496,10 @@ index f0f3f94b811..b7d18771314 100644 spark.internalCreateDataFrame(withoutFilters.execute(), schema) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -index 720b13b812e..e3ac2cebc6e 100644 +index 720b13b812e..388dad6b4f3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -@@ -98,6 +98,20 @@ trait SharedSparkSessionBase +@@ -98,6 +98,25 @@ trait SharedSparkSessionBase // this rule may potentially block testing of other optimization rules such as // ConstantPropagation etc. .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName) @@ -4078,6 +4516,11 @@ index 720b13b812e..e3ac2cebc6e 100644 + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") + ++ // CometDriverPlugin installs Comet's cache serializer when ++ // spark.comet.exec.inMemoryCache.enabled is on, as it is by default. These sessions do not ++ // load the plugin, so install it here. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } conf.set( StaticSQLConf.WAREHOUSE_PATH, @@ -4200,10 +4643,10 @@ index b67370f6eb9..746b3974b29 100644 override def beforeEach(): Unit = { super.beforeEach() diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala -index a394d0b7393..7056c74759c 100644 +index a394d0b7393..9c626b05ab2 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala -@@ -53,24 +53,34 @@ object TestHive +@@ -53,24 +53,38 @@ object TestHive new SparkContext( System.getProperty("spark.sql.test.master", "local[1]"), "TestSQLContext", @@ -4248,6 +4691,10 @@ index a394d0b7393..7056c74759c 100644 + .set("spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") ++ ++ // As in SharedSparkSession: what CometDriverPlugin would install. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } + + conf diff --git a/dev/diffs/4.1.3.diff b/dev/diffs/4.1.3.diff index c99c1c242d2..0b814aa1b10 100644 --- a/dev/diffs/4.1.3.diff +++ b/dev/diffs/4.1.3.diff @@ -78,7 +78,7 @@ index 370bfa4397f..ccb4398de58 100644 org.apache.datasketches diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala -index 01aa71047fe..cde29aabc83 100644 +index 01aa71047fe..06913c0656b 100644 --- a/project/SparkBuild.scala +++ b/project/SparkBuild.scala @@ -443,9 +443,11 @@ object SparkBuild extends PomBuild { @@ -93,7 +93,7 @@ index 01aa71047fe..cde29aabc83 100644 enable(HiveThriftServer.settings)(hiveThriftServer) -@@ -1347,6 +1349,25 @@ object SQL { +@@ -1347,6 +1349,35 @@ object SQL { } } @@ -101,16 +101,26 @@ index 01aa71047fe..cde29aabc83 100644 + * Comet disables itself when its shuffle manager is not registered. When Comet is enabled, make + * that shuffle manager the default for every test SparkConf in the modules that have Comet on + * their classpath, so that suites which build their own SparkSession or SparkContext run Comet -+ * too, not just the ones that go through SharedSparkSession or TestHive. ++ * too, not just the ones that go through SharedSparkSession or TestHive. Comet's cache serializer ++ * is made the default the same way. Spark keeps the first cache serializer it loads for the rest ++ * of the JVM, so a suite that caches through a session of its own would otherwise leave the ++ * suites after it in Spark's cache format. + */ +object CometTestSettings { + private val cometEnabled = + sys.env.get("ENABLE_COMET").forall(v => v == "1" || v.equalsIgnoreCase("true")) + ++ private val defaults = Seq( ++ "spark.shuffle.manager" -> ++ "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager", ++ "spark.sql.cache.serializer" -> ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") ++ + lazy val settings: Seq[Setting[_]] = -+ if (cometEnabled && !sys.props.contains("spark.shuffle.manager")) { -+ Seq((Test / javaOptions) += -+ "-Dspark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") ++ if (cometEnabled) { ++ defaults.filterNot { case (key, _) => sys.props.contains(key) }.map { ++ case (key, value) => (Test / javaOptions) += s"-D$key=$value" ++ } + } else { + Seq.empty + } @@ -400,11 +410,41 @@ index 21a3ce1e122..f4762ab98f0 100644 SET spark.sql.ansi.enabled = false; -- In COMPENSATION views get invalidated if the type can't cast +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CacheTableInKryoSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CacheTableInKryoSuite.scala +index 26d8f750f6e..c888f8e0844 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/CacheTableInKryoSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/CacheTableInKryoSuite.scala +@@ -32,9 +32,16 @@ class CacheTableInKryoSuite extends QueryTest + with SharedSparkSession { + + override def sparkConf: SparkConf = { +- super.sparkConf ++ val conf = super.sparkConf + .set("spark.kryo.registrationRequired", "true") + .set("spark.serializer", "org.apache.spark.serializer.KryoSerializer") ++ // Comet's cache format needs its classes registered too, which is what Comet asks of any ++ // application that sets registrationRequired. ++ if (isCometEnabled) { ++ conf.set("spark.kryo.registrator", "org.apache.comet.CometKryoRegistrator") ++ } else { ++ conf ++ } + } + + test("SPARK-51777: sql.columnar.* classes registered in KryoSerializer") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala -index 0d807aeae4d..6d7744e771b 100644 +index 0d807aeae4d..b5a00b75580 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala -@@ -49,7 +49,7 @@ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEProp +@@ -37,6 +37,7 @@ import org.apache.spark.sql.catalyst.analysis.TempTableAlreadyExistsException + import org.apache.spark.sql.catalyst.expressions.SubqueryExpression + import org.apache.spark.sql.catalyst.plans.logical.{BROADCAST, Join, JoinStrategyHint, SHUFFLE_HASH} + import org.apache.spark.sql.catalyst.util.DateTimeConstants ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.connector.catalog.BasicInMemoryTableCatalog + import org.apache.spark.sql.connector.catalog.CatalogPlugin + import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.CatalogHelper +@@ -49,7 +50,7 @@ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEProp import org.apache.spark.sql.execution.columnar._ import org.apache.spark.sql.execution.command.CommandUtils import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation @@ -413,7 +453,28 @@ index 0d807aeae4d..6d7744e771b 100644 import org.apache.spark.sql.execution.ui.SparkListenerSQLAdaptiveExecutionUpdate import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -@@ -534,7 +534,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils +@@ -128,6 +129,10 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + case inMemoryTable @ InMemoryTableScanExec(_, _, relation) => + getNumInMemoryTablesRecursively(relation.cachedPlan) + + getNumInMemoryTablesInSubquery(inMemoryTable) + 1 ++ // Comet's cache scan keeps the predicates pushed into it on its original plan, out of reach ++ // of subquery reuse. The filter above it evaluates the same subqueries, and counts them. ++ case cometScan: CometInMemoryTableScanExec => ++ getNumInMemoryTablesRecursively(cometScan.originalPlan.relation.cachedPlan) + 1 + case p => + getNumInMemoryTablesInSubquery(p) + }.sum +@@ -408,7 +413,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + assert(isExpectStorageLevel(rddId, Disk)) + } + +- test("InMemoryRelation statistics") { ++ test("InMemoryRelation statistics", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + sql("CACHE TABLE testData") + spark.table("testData").queryExecution.withCachedData.collect { + case cached: InMemoryRelation => +@@ -534,7 +540,8 @@ class CachedTableSuite extends QueryTest with SQLTestUtils df.collect() } assert( @@ -423,7 +484,19 @@ index 0d807aeae4d..6d7744e771b 100644 } test("A cached table preserves the partitioning and ordering of its cached SparkPlan") { -@@ -1673,9 +1674,18 @@ class CachedTableSuite extends QueryTest with SQLTestUtils +@@ -1602,7 +1609,10 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + checkAnswer(spark.table(tableName), Row(LocalDateTime.parse("2021-01-01T00:00:00"))) + spark.table(tableName).queryExecution.withCachedData.collect { + case cached: InMemoryRelation => +- assert(cached.stats.sizeInBytes === 8) ++ // Comet's cache format reports Arrow buffer sizes. ++ if (!isCometEnabled) { ++ assert(cached.stats.sizeInBytes === 8) ++ } + } + sql(s"UNCACHE TABLE $tableName") + } +@@ -1673,9 +1683,19 @@ class CachedTableSuite extends QueryTest with SQLTestUtils _.nodeName.contains("TableCacheQueryStage")) val aqeNode = findNodeInSparkPlanInfo(inMemoryScanNode.get, _.nodeName.contains("AdaptiveSparkPlan")) @@ -431,20 +504,45 @@ index 0d807aeae4d..6d7744e771b 100644 - _.nodeName.contains("ResultQueryStage")) - aqePlanRoot.get.children.head.nodeName == "AQEShuffleRead" + // Spark 4.0 wraps results in ResultQueryStage. The coalescing indicator is AQEShuffleRead -+ // as the direct child of InputAdapter. ++ // as its direct child, as in Spark, when the table is cached in Comet's columnar format, ++ // or as the direct child of InputAdapter when Comet's plan is cached in Spark's format. + // AdaptiveSparkPlan -> ResultQueryStage -> WholestageCodegen -> + // CometColumnarToRow -> InputAdapter -> AQEShuffleRead (if coalesced) + val resultStage = aqeNode.get.children.head // ResultQueryStage -+ val wsc = resultStage.children.head // WholeStageCodegen ++ val wsc = resultStage.children.head // WholeStageCodegen or AQEShuffleRead + val c2r = wsc.children.head // ColumnarToRow or CometColumnarToRow + val inputAdapter = c2r.children.head // InputAdapter -+ resultStage.nodeName == "ResultQueryStage" && ++ resultStage.nodeName == "ResultQueryStage" && (wsc.nodeName == "AQEShuffleRead" || + wsc.nodeName.startsWith("WholeStageCodegen") && // could be "WholeStageCodegen (1)" + (c2r.nodeName == "CometColumnarToRow" || c2r.nodeName == "ColumnarToRow") && -+ inputAdapter.children.head.nodeName == "AQEShuffleRead" ++ inputAdapter.children.head.nodeName == "AQEShuffleRead") } withTempView("t0", "t1", "t2") { +@@ -1764,6 +1784,7 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + val cached = spark.table("t") + val tableCache = collect(cached.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c.originalPlan + } + if (expected == StorageLevel.NONE) { + assert(tableCache.isEmpty) +@@ -2630,6 +2651,7 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + + val inMemoryTableScan = collect(df.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c.originalPlan + } + assert(inMemoryTableScan.size == 1) + checkAnswer(df, Row(5) :: Nil) +@@ -2657,6 +2679,7 @@ class CachedTableSuite extends QueryTest with SQLTestUtils + + val subqueryInMemoryTableScan = collect(cteInSubquery.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c.originalPlan + } + assert(subqueryInMemoryTableScan.size == 1) + checkAnswer(cteInSubquery, Row(1) :: Nil) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala index bfe15b33768..df8fb4c0a35 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -512,7 +610,7 @@ index ed182322aec..1ae6afa686a 100644 spark.range(100).write.saveAsTable(s"$dbName.$table2Name") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala -index 93ff7becaec..87537a25b3b 100644 +index 93ff7becaec..78b02628d42 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala @@ -23,10 +23,11 @@ import java.util.Locale @@ -528,7 +626,24 @@ index 93ff7becaec..87537a25b3b 100644 import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.{ExamplePoint, ExamplePointUDT, SharedSparkSession, SQLTestData} -@@ -1519,11 +1520,12 @@ class DataFrameSetOperationsSuite extends QueryTest +@@ -1497,9 +1498,13 @@ class DataFrameSetOperationsSuite extends QueryTest + + val union = df1.union(df2) + union.collect() +- checkIfColumnar(union.queryExecution.executedPlan, +- _.isInstanceOf[InMemoryTableScanExec], supported, 2) +- checkIfColumnar(union.queryExecution.executedPlan, _.isInstanceOf[UnionExec], supported, 1) ++ // Comet replaces the cache scans and the union with its own operators. ++ if (!isCometEnabled) { ++ checkIfColumnar(union.queryExecution.executedPlan, ++ _.isInstanceOf[InMemoryTableScanExec], supported, 2) ++ checkIfColumnar(union.queryExecution.executedPlan, ++ _.isInstanceOf[UnionExec], supported, 1) ++ } + checkAnswer(union, Row(1) :: Row(2) :: Row(3) :: Row(4) :: Row(5) :: Row(6) :: Nil) + + val nonColumnarUnion = df1.union(Seq(7, 8, 9).toDF("k")) +@@ -1519,11 +1524,12 @@ class DataFrameSetOperationsSuite extends QueryTest val union = df1.repartition($"a").union(df2.repartition($"a")) val unionExec = union.queryExecution.executedPlan.collect { case u: UnionExec => u @@ -542,7 +657,7 @@ index 93ff7becaec..87537a25b3b 100644 } assert(shuffle.size == 1) -@@ -1554,11 +1556,12 @@ class DataFrameSetOperationsSuite extends QueryTest +@@ -1554,11 +1560,12 @@ class DataFrameSetOperationsSuite extends QueryTest val union = df1.repartition($"a").union(df2.repartition($"d")) val unionExec = union.queryExecution.executedPlan.collect { case u: UnionExec => u @@ -556,7 +671,7 @@ index 93ff7becaec..87537a25b3b 100644 } assert(shuffle.size == 1) -@@ -1573,10 +1576,10 @@ class DataFrameSetOperationsSuite extends QueryTest +@@ -1573,10 +1580,10 @@ class DataFrameSetOperationsSuite extends QueryTest // Avoid unnecessary shuffle if union output partitioning is enabled val shuffledUnion = union.repartition($"a") val shuffleNumBefore = union.queryExecution.executedPlan.collect { @@ -569,7 +684,7 @@ index 93ff7becaec..87537a25b3b 100644 } if (enabled) { -@@ -1605,6 +1608,7 @@ class DataFrameSetOperationsSuite extends QueryTest +@@ -1605,6 +1612,7 @@ class DataFrameSetOperationsSuite extends QueryTest val union = df1.repartitionByRange($"a").union(df2.repartitionByRange($"d")) val unionExec = union.queryExecution.executedPlan.collect { case u: UnionExec => u @@ -710,6 +825,99 @@ index 4a070becfa6..61d515d127e 100644 val df = Seq((1, "1"), (2, "2"), (1, "3"), (2, "4")).toDF("key", "value") val window = Window.partitionBy($"key").orderBy($"value") +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +index 627811eaecf..7269aeab949 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +@@ -22,7 +22,8 @@ import java.time.LocalTime + import org.scalatest.concurrent.TimeLimits + import org.scalatest.time.SpanSugar._ + +-import org.apache.spark.sql.execution.ColumnarToRowExec ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec ++import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.columnar.{InMemoryRelation, InMemoryTableScanExec} + import org.apache.spark.sql.functions._ +@@ -39,6 +40,10 @@ class DatasetCacheSuite extends QueryTest + with AdaptiveSparkPlanHelper { + import testImplicits._ + ++ // A scan of a cached relation, Spark's or Comet's. ++ private def isCacheScan(plan: SparkPlan): Boolean = ++ plan.isInstanceOf[InMemoryTableScanExec] || plan.isInstanceOf[CometInMemoryTableScanExec] ++ + /** + * Asserts that a cached [[Dataset]] will be built using the given number of other cached results. + */ +@@ -46,7 +51,7 @@ class DatasetCacheSuite extends QueryTest + val plan = df.queryExecution.withCachedData + assert(plan.isInstanceOf[InMemoryRelation]) + val internalPlan = plan.asInstanceOf[InMemoryRelation].cacheBuilder.cachedPlan +- assert(find(internalPlan)(_.isInstanceOf[InMemoryTableScanExec]).size ++ assert(find(internalPlan)(isCacheScan).size + == numOfCachesDependedUpon) + } + +@@ -256,7 +261,7 @@ class DatasetCacheSuite extends QueryTest + case i: InMemoryRelation => i.cacheBuilder.cachedPlan + } + assert(df2LimitInnerPlan.isDefined && +- !df2LimitInnerPlan.get.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ !df2LimitInnerPlan.get.exists(isCacheScan)) + } + + test("SPARK-27739 Save stats from optimized plan") { +@@ -289,14 +294,14 @@ class DatasetCacheSuite extends QueryTest + val unionDf = df1.union(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.union(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("positive: union by name") { + val unionDf = df1.unionByName(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.unionByName(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by position") { +@@ -304,7 +309,7 @@ class DatasetCacheSuite extends QueryTest + unionDf.cache() + val finalDf = unionDf.union(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by name") { +@@ -312,7 +317,7 @@ class DatasetCacheSuite extends QueryTest + unionDf.cache() + val finalDf = unionDf.unionByName(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + } + } +@@ -324,12 +329,12 @@ class DatasetCacheSuite extends QueryTest + df1.cache() + // This is exactly the same as df1. + val df2 = spark.range(5).select(struct($"id".as("name", metadata))) +- assert(df2.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(df2.queryExecution.executedPlan.exists(isCacheScan)) + + val metadata2 = Metadata.fromJson("""{"k2": "v2"}""") + // Same with df1 except for the Alias metadata + val df3 = spark.range(5).select(struct($"id".as("name", metadata2))) +- assert(!df3.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!df3.queryExecution.executedPlan.exists(isCacheScan)) + } + + test("SPARK-53418: Handle TimeType in ColumnAccessor") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala index 6df8d66ee7f..35e270c7241 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala @@ -733,10 +941,10 @@ index 6df8d66ee7f..35e270c7241 100644 assert(exchanges.size == 2) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -index e1a2fd33c7c..632f4b695df 100644 +index e1a2fd33c7c..ac0682e0e1a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -@@ -22,6 +22,7 @@ import org.scalatest.GivenWhenThen +@@ -22,9 +22,11 @@ import org.scalatest.GivenWhenThen import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression} import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode._ import org.apache.spark.sql.catalyst.plans.ExistenceJoin @@ -744,7 +952,11 @@ index e1a2fd33c7c..632f4b695df 100644 import org.apache.spark.sql.connector.catalog.{InMemoryTableCatalog, InMemoryTableWithV2FilterCatalog} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive._ -@@ -193,6 +194,7 @@ abstract class DynamicPartitionPruningSuiteBase ++import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec + import org.apache.spark.sql.execution.datasources.v2.BatchScanExec + import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ReusedExchangeExec} + import org.apache.spark.sql.execution.joins.BroadcastHashJoinExec +@@ -193,6 +195,7 @@ abstract class DynamicPartitionPruningSuiteBase } val subqueryBroadcast = dpExprs.collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b @@ -752,7 +964,19 @@ index e1a2fd33c7c..632f4b695df 100644 } val hasFilter = if (withSubquery) "Should" else "Shouldn't" -@@ -247,6 +249,8 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -229,7 +232,10 @@ abstract class DynamicPartitionPruningSuiteBase + } + + val isMainQueryAdaptive = plan.isInstanceOf[AdaptiveSparkPlanExec] +- subqueriesAll(plan).filterNot(subqueryBroadcast.contains).foreach { s => ++ // Comet's cache scan exposes Spark's own scan as its one subquery, for Spark's UI to draw ++ // the cached plan below it, but the query never runs it as a subquery. ++ subqueriesAll(plan).filterNot(subqueryBroadcast.contains) ++ .filterNot(_.isInstanceOf[InMemoryTableScanExec]).foreach { s => + val subquery = s match { + case r: ReusedSubqueryExec => r.child + case o => o +@@ -247,6 +253,8 @@ abstract class DynamicPartitionPruningSuiteBase val buf = collectDynamicPruningExpressions(df.queryExecution.executedPlan).collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b.indices.map(idx => b.buildKeys(idx)) @@ -761,7 +985,7 @@ index e1a2fd33c7c..632f4b695df 100644 } assert(buf.distinct.size == n) } -@@ -262,6 +266,12 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -262,6 +270,12 @@ abstract class DynamicPartitionPruningSuiteBase case s: BatchScanExec => s.runtimeFilters.collect { case d: DynamicPruningExpression => d.child } @@ -774,7 +998,7 @@ index e1a2fd33c7c..632f4b695df 100644 case _ => Nil } } -@@ -1204,10 +1214,16 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1204,10 +1218,16 @@ abstract class DynamicPartitionPruningSuiteBase val plan = df.queryExecution.executedPlan val countSubqueryBroadcasts = @@ -793,7 +1017,7 @@ index e1a2fd33c7c..632f4b695df 100644 assert(countSubqueryBroadcasts == 1) assert(countReusedSubqueryBroadcasts == 1) -@@ -1579,6 +1595,7 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1579,6 +1599,7 @@ abstract class DynamicPartitionPruningSuiteBase val subqueryBroadcastExecs = collectWithSubqueries(df.queryExecution.executedPlan) { case s: SubqueryBroadcastExec => s @@ -801,7 +1025,7 @@ index e1a2fd33c7c..632f4b695df 100644 } assert(subqueryBroadcastExecs.size === 1) subqueryBroadcastExecs.foreach { subqueryBroadcastExec => -@@ -1731,6 +1748,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat +@@ -1731,6 +1752,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat case s: BatchScanExec => // we use f1 col for v2 tables due to schema pruning s.output.exists(_.exists(_.argString(maxFields = 100).contains("f1"))) @@ -813,7 +1037,7 @@ index e1a2fd33c7c..632f4b695df 100644 } assert(scanOption.isDefined) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala -index 4c62c47971a..ecc7ed21546 100644 +index 4c62c47971a..fe788debf70 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala @@ -268,7 +268,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite @@ -826,6 +1050,15 @@ index 4c62c47971a..ecc7ed21546 100644 withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true") { withTempView("df") { val df1 = spark.range(1, 100) +@@ -339,7 +340,7 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite + + val sqlText = "EXPLAIN FORMATTED SELECT (SELECT min(id) FROM df) as v" + val expected_pattern = +- "Subquery:1 Hosting operator id = 2 Hosting Expression = Subquery subquery#x" ++ "Subquery:1 Hosting operator id = \\d+ Hosting Expression = Subquery subquery#x" + + withNormalizedExplain(sqlText) { normalizedOutput => + assert(expected_pattern.r.findAllMatchIn(normalizedOutput).length == 1) @@ -471,7 +472,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite } } @@ -1313,10 +1546,18 @@ index cb9d0909554..084d6515e8b 100644 } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala -index 74cdee49e55..f7452c9abb7 100644 +index 74cdee49e55..da1e9bd966a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala -@@ -1521,7 +1521,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -36,6 +36,7 @@ import org.apache.spark.sql.catalyst.optimizer.{ConvertToLocalRelation, NestedCo + import org.apache.spark.sql.catalyst.parser.ParseException + import org.apache.spark.sql.catalyst.plans.logical.{LocalLimit, Project, RepartitionByExpression, Sort} + import org.apache.spark.sql.connector.catalog.CatalogManager.SESSION_CATALOG_NAME ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.{CommandResultExec, OneRowRelationExec, UnionExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.aggregate._ +@@ -1521,7 +1522,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark checkAnswer(sql("select -0.001"), Row(BigDecimal("-0.001"))) } @@ -1326,7 +1567,7 @@ index 74cdee49e55..f7452c9abb7 100644 AccumulatorSuite.verifyPeakExecutionMemorySet(sparkContext, "external sort") { sql("SELECT * FROM testData2 ORDER BY a ASC, b ASC").collect() } -@@ -1982,8 +1983,15 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark +@@ -1982,8 +1984,15 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark countAcc.add(1) x }) @@ -1343,6 +1584,24 @@ index 74cdee49e55..f7452c9abb7 100644 verifyCallCount( df.selectExpr("testUdf(a + 1) + testUdf(1 + a)", "testUdf(a + 1)"), Row(4, 2), 1) +@@ -3275,6 +3284,7 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark + + val inMemoryTableScan = collect(queryDf.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c + } + assert(inMemoryTableScan.size == 2) + checkAnswer(queryDf, Row(0, 1) :: Row(1, 2) :: Row(2, 3) :: Row(3, 4) :: Row(4, 5) :: Nil) +@@ -4741,7 +4751,8 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark + ) + } + +- test("SPARK-47939: Explain should work with parameterized queries") { ++ test("SPARK-47939: Explain should work with parameterized queries", ++ IgnoreComet("Comet explain output is different")) { + def checkQueryPlan(df: DataFrame, plan: String): Unit = assert( + df.collect() + .map(_.getString(0)) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala index 23f0144dcec..40d536bb23a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala @@ -1388,7 +1647,7 @@ index 23f0144dcec..40d536bb23a 100644 } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala -index 66826a9ca76..a330d142eb1 100644 +index 66826a9ca76..db56c426dc0 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala @@ -196,7 +196,15 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt @@ -1407,7 +1666,15 @@ index 66826a9ca76..a330d142eb1 100644 val extensions = create { extensions => // inject rule that will run during AQE query stage preparation and will add custom tags // to the plan -@@ -290,6 +298,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -239,6 +247,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt + df.select("i").filter($"i" > 1).cache() + assert(find(df.filter($"i" > 1).select("i").queryExecution.executedPlan) { + case _: org.apache.spark.sql.execution.columnar.InMemoryTableScanExec => true ++ case _: org.apache.spark.sql.comet.CometInMemoryTableScanExec => true + case _ => false + }.isDefined) + } +@@ -290,6 +299,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } private def testInjectColumnar(enableAQE: Boolean): Unit = { @@ -1415,7 +1682,7 @@ index 66826a9ca76..a330d142eb1 100644 def collectPlanSteps(plan: SparkPlan): Seq[Int] = plan match { case a: AdaptiveSparkPlanExec => assert(a.toString.startsWith("AdaptiveSparkPlan isFinalPlan=true")) -@@ -345,6 +354,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -345,6 +355,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } test("reset column vectors") { @@ -1423,7 +1690,7 @@ index 66826a9ca76..a330d142eb1 100644 val session = SparkSession.builder() .master("local[1]") .config(COLUMN_BATCH_SIZE.key, 2) -@@ -513,6 +523,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt +@@ -513,6 +524,7 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt } test("SPARK-38697: Extend SparkSessionExtensions to inject rules into AQE Optimizer") { @@ -1446,6 +1713,33 @@ index d7b2511eac2..d5f5b940b94 100644 sc = new SparkContext("local[2]", "test") val session = classic.SparkSession.builder().sparkContext(sc).getOrCreate() import session.implicits._ +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +index 5222d5ce266..6f8f8bec3d4 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +@@ -847,14 +847,19 @@ class StatisticsCollectionSuite extends StatisticsCollectionTestBase with Shared + checkTableStats("t1", hasSizeInBytes = true, expectedRowCounts = None) + checkTableStats("t2", hasSizeInBytes = true, expectedRowCounts = None) + assert(getCatalogTable("v1").stats.isEmpty) +- checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ // Comet's cache format reports Arrow buffer sizes for a materialized cached view. ++ if (!isCometEnabled) { ++ checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ } + checkOptimizedPlanStats(spark.table("v2"), 1, None, Seq.empty) + + sql("ANALYZE TABLES COMPUTE STATISTICS") + checkTableStats("t1", hasSizeInBytes = true, expectedRowCounts = Some(10)) + checkTableStats("t2", hasSizeInBytes = true, expectedRowCounts = Some(20)) +- checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) +- checkOptimizedPlanStats(spark.table("v2"), 4, Some(1), Seq.empty) ++ if (!isCometEnabled) { ++ checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ checkOptimizedPlanStats(spark.table("v2"), 4, Some(1), Seq.empty) ++ } + } + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala index 7bfc8cf4fa6..4bd387801db 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala @@ -1645,7 +1939,7 @@ index 8a0e2c29653..d276a51cbc6 100644 } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala -index 8f7a68bcbe6..88dbe1793c9 100644 +index 8f7a68bcbe6..c09c5d74309 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala @@ -26,6 +26,8 @@ import org.apache.spark.sql.{AnalysisException, Row} @@ -1653,7 +1947,7 @@ index 8f7a68bcbe6..88dbe1793c9 100644 import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.util.CollationFactory +import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashJoinExec, CometSortMergeJoinExec} -+import org.apache.spark.sql.comet.CometHashAggregateExec ++import org.apache.spark.sql.comet.{CometHashAggregateExec, CometInMemoryTableScanExec} import org.apache.spark.sql.connector.{DatasourceV2SQLBase, FakeV2ProviderWithCustomSchema} import org.apache.spark.sql.connector.catalog.{CatalogV2Util, Identifier, InMemoryTable} import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.CatalogHelper @@ -1712,6 +2006,15 @@ index 8f7a68bcbe6..88dbe1793c9 100644 }.head.isInstanceOf[ArrayTransform]) } } +@@ -1948,7 +1960,7 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper { + // Checks in-memory fetching code path. + val all = sql("SELECT col FROM tbl") + assert(all.queryExecution.executedPlan.collectFirst { +- case _: InMemoryTableScanExec => true ++ case _: InMemoryTableScanExec | _: CometInMemoryTableScanExec => true + }.nonEmpty) + checkAnswer(all, Row("a")) + // Checks column stats code path. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala index a09b7e0827c..ffc29f764bc 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala @@ -1894,7 +2197,7 @@ index fcecaf25d4c..e5a511022cc 100644 // Fail to read ancient datetime values. withSQLConf(SQLConf.PARQUET_REBASE_MODE_IN_READ.key -> EXCEPTION.toString) { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala -index 28762f01d7a..6cdc3953bf3 100644 +index 28762f01d7a..d19f7ff9d34 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala @@ -56,6 +56,11 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper @@ -1954,23 +2257,6 @@ index 28762f01d7a..6cdc3953bf3 100644 } } } -@@ -475,10 +486,12 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper - // Shuffle partition coalescing of the join is performed independent of the non-grouping - // aggregate on the other side of the union. - val finalPlan = stripAQEPlan(resultDf.queryExecution.executedPlan) -- assert( -- finalPlan.collect { -- case r @ CoalescedShuffleRead() => r -- }.size == 2) -+ if (!isCometEnabled) { -+ assert( -+ finalPlan.collect { -+ case r @ CoalescedShuffleRead() => r -+ }.size == 2) -+ } - } - withSparkSession(test, 100, None) - } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala index 418ca3430bb..eb8267192f8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala @@ -2521,7 +2807,7 @@ index a3cfdc5a240..3793b6191bf 100644 }) checkAnswer(distinctWithId, Seq(Row(1, 0), Row(1, 0))) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala -index 188a28ff1c0..2717d45f3d0 100644 +index 188a28ff1c0..d1d8714cd32 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala @@ -27,12 +27,14 @@ import org.apache.spark.SparkException @@ -3006,7 +3292,18 @@ index 188a28ff1c0..2717d45f3d0 100644 }.isEmpty) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -3229,7 +3282,8 @@ class AdaptiveQueryExecSuite +@@ -3223,13 +3276,18 @@ class AdaptiveQueryExecSuite + + val (_, adaptive) = runAdaptiveAndVerifyResult( + "SELECT * FROM t1 WHERE c1 < (SELECT c2 FROM t2)") +- assert(findReusedSubquery(adaptive).size == 1) ++ // Comet's cache scan does not plan the subqueries in its pruning predicates, so only the ++ // filter above it runs this one and there is nothing to reuse. ++ if (!isCometEnabled) { ++ assert(findReusedSubquery(adaptive).size == 1) ++ } + } + } } } @@ -3016,7 +3313,7 @@ index 188a28ff1c0..2717d45f3d0 100644 val emptyDf = spark.range(1).where("false") val aggDf1 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df1")) val aggDf2 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df2")) -@@ -3319,7 +3373,8 @@ class AdaptiveQueryExecSuite +@@ -3319,7 +3377,8 @@ class AdaptiveQueryExecSuite val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec]) @@ -3027,10 +3324,29 @@ index 188a28ff1c0..2717d45f3d0 100644 assert(o1.semanticEquals(o2), "Different output column order after AQE optimization") } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala -index 47b935a2880..3fdeab3113c 100644 +index 47b935a2880..65ee66c1975 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala -@@ -230,9 +230,21 @@ class CachedBatchSerializerNoUnwrapSuite extends QueryTest +@@ -212,6 +212,18 @@ class CachedBatchSerializerNoUnwrapSuite extends QueryTest + classOf[DefaultCachedBatchSerializerNoUnwrap].getName) + } + ++ // Spark keeps the first cache serializer it loads for the rest of the JVM, which the suites ++ // before this one set to Comet's. Clear it on both sides, as CachedBatchSerializerSuite does. ++ protected override def beforeAll(): Unit = { ++ super.beforeAll() ++ clearSerializer() ++ } ++ ++ protected override def afterAll(): Unit = { ++ clearSerializer() ++ super.afterAll() ++ } ++ + test("Do not unwrap ColumnarToRowExec") { + withTempPath { workDir => + val workDirPath = workDir.getAbsolutePath +@@ -230,9 +242,21 @@ class CachedBatchSerializerNoUnwrapSuite extends QueryTest assert(cachedPlans.length == 2) cachedPlans.foreach { cachedPlan => @@ -3055,6 +3371,147 @@ index 47b935a2880..3fdeab3113c 100644 } } } +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +index 4f07d3d1c03..9610c703610 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +@@ -28,6 +28,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, + import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning + import org.apache.spark.sql.classic.DataFrame + import org.apache.spark.sql.columnar.CachedBatch ++import org.apache.spark.sql.comet.{CometFilterExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{FilterExec, InputAdapter, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions._ +@@ -510,11 +511,17 @@ class InMemoryColumnarQuerySuite extends QueryTest + val planBeforeFilter = collect(df2.queryExecution.executedPlan) { + case f: FilterExec => f.child + case WholeStageCodegenExec(FilterExec(_, i: InputAdapter)) => i.child ++ case f: CometFilterExec => f.child + } +- assert(planBeforeFilter.head.isInstanceOf[InMemoryTableScanExec]) +- + val execPlan = planBeforeFilter.head +- assert(execPlan.executeCollectPublic().length == 0) ++ execPlan match { ++ // Comet's cache scan is columnar only, so count the rows of the batches it emits. ++ case c: CometInMemoryTableScanExec => ++ assert(c.executeColumnar().map(_.numRows().toLong).collect().sum == 0) ++ case _ => ++ assert(execPlan.isInstanceOf[InMemoryTableScanExec]) ++ assert(execPlan.executeCollectPublic().length == 0) ++ } + } + + test("SPARK-25727 - otherCopyArgs in InMemoryRelation does not include outputOrdering") { +@@ -544,7 +551,10 @@ class InMemoryColumnarQuerySuite extends QueryTest + + // InMemoryRelation's stats is updated after materializing RDD + dfFromFile.collect() +- assert(inMemoryRelation.computeStats().sizeInBytes === 16) ++ // Comet's cache format reports Arrow buffer sizes. ++ if (!isCometEnabled) { ++ assert(inMemoryRelation.computeStats().sizeInBytes === 16) ++ } + + // test of catalog table + val dfFromTable = spark.catalog.createTable("table1", workDirPath).cache() +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +index 2c73622739a..5d0efeb263a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +@@ -18,6 +18,7 @@ + package org.apache.spark.sql.execution.columnar + + import org.apache.spark.SparkFunSuite ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.SparkPlan + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions.expr +@@ -37,11 +38,15 @@ class InMemoryRelationSuite extends SparkFunSuite + test("SPARK-47177: Cached SQL plan do not display final AQE plan in explain string") { + def findIMRInnerChild(p: SparkPlan): SparkPlan = { + val tableCache = find(p) { +- case _: InMemoryTableScanExec => true ++ case _: InMemoryTableScanExec | _: CometInMemoryTableScanExec => true + case _ => false + } + assert(tableCache.isDefined) +- tableCache.get.asInstanceOf[InMemoryTableScanExec].relation.innerChildren.head ++ val scan = tableCache.get match { ++ case c: CometInMemoryTableScanExec => c.originalPlan ++ case s => s.asInstanceOf[InMemoryTableScanExec] ++ } ++ scan.relation.innerChildren.head + } + + val d1 = spark.range(1).withColumn("key", expr("id % 100")) +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +index 88ff51d0ff4..7d00fca3b0a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +@@ -17,6 +17,7 @@ + + package org.apache.spark.sql.execution.columnar + ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.internal.SQLConf + import org.apache.spark.sql.test.SharedSparkSession +@@ -181,11 +182,16 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + val result = df.collect().map(_(0)).toArray + assert(result.length === 1) + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- assert(readPartitions === 5) +- assert(readBatches === 10) ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan has none of these test-only accumulators, so there is nothing to count. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ assert(scans.head.readPartitions.value === 5) ++ assert(scans.head.readBatches.value === 10) ++ } + } + + def checkBatchPruning( +@@ -202,14 +208,23 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + df.collect().map(_(0)).toArray + } + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- +- assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution") +- assert( +- readPartitions === expectedReadPartitions, +- s"Wrong number of read partitions: $queryExecution") ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan prunes on the same statistics, but has none of these test-only ++ // accumulators to read, so only the answer above is checked for it. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ val readPartitions = scans.head.readPartitions.value ++ val readBatches = scans.head.readBatches.value ++ assert( ++ readBatches === expectedReadBatches, ++ s"Wrong number of read batches: $queryExecution") ++ assert( ++ readPartitions === expectedReadPartitions, ++ s"Wrong number of read partitions: $queryExecution") ++ } + } + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala index 269990d7d14..140ee4112b1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala @@ -3542,10 +3999,10 @@ index b8f3ea3c6f3..bbd44221288 100644 val workDirPath = workDir.getAbsolutePath val input = spark.range(5).toDF("id") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -index 9bd858608cb..2682ba53513 100644 +index 9bd858608cb..896d33a8f0e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -@@ -21,7 +21,7 @@ import scala.reflect.ClassTag +@@ -21,13 +21,14 @@ import scala.reflect.ClassTag import org.apache.spark.AccumulatorSuite import org.apache.spark.internal.config.EXECUTOR_MEMORY @@ -3554,7 +4011,29 @@ index 9bd858608cb..2682ba53513 100644 import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BitwiseAnd, BitwiseOr, Cast, Expression, Literal, ShiftLeft} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.Inner -@@ -488,7 +488,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + import org.apache.spark.sql.catalyst.plans.logical.BROADCAST + import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, PartitioningCollection} + import org.apache.spark.sql.classic.{Dataset, SparkSession} ++import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{DummySparkPlan, SparkPlan, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} + import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec +@@ -158,12 +159,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) +@@ -488,7 +491,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils } } @@ -3564,6 +4043,21 @@ index 9bd858608cb..2682ba53513 100644 withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") { val t1 = (0 until 100).map(i => (i % 5, i % 13)).toDF("i1", "j1") val t2 = (0 until 100).map(i => (i % 5, i % 14)).toDF("i2", "j2") +@@ -656,12 +660,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest with SQLTestUtils + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala index f2e9121d566..2c9f517034f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala @@ -4371,10 +4865,10 @@ index f0f3f94b811..b7d18771314 100644 spark.internalCreateDataFrame(withoutFilters.execute(), schema) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -index 720b13b812e..e3ac2cebc6e 100644 +index 720b13b812e..388dad6b4f3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -@@ -98,6 +98,20 @@ trait SharedSparkSessionBase +@@ -98,6 +98,25 @@ trait SharedSparkSessionBase // this rule may potentially block testing of other optimization rules such as // ConstantPropagation etc. .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName) @@ -4391,6 +4885,11 @@ index 720b13b812e..e3ac2cebc6e 100644 + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") + ++ // CometDriverPlugin installs Comet's cache serializer when ++ // spark.comet.exec.inMemoryCache.enabled is on, as it is by default. These sessions do not ++ // load the plugin, so install it here. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } conf.set( StaticSQLConf.WAREHOUSE_PATH, @@ -4513,10 +5012,10 @@ index b67370f6eb9..746b3974b29 100644 override def beforeEach(): Unit = { super.beforeEach() diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala -index a394d0b7393..7056c74759c 100644 +index a394d0b7393..9c626b05ab2 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala -@@ -53,24 +53,34 @@ object TestHive +@@ -53,24 +53,38 @@ object TestHive new SparkContext( System.getProperty("spark.sql.test.master", "local[1]"), "TestSQLContext", @@ -4561,6 +5060,10 @@ index a394d0b7393..7056c74759c 100644 + .set("spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") ++ ++ // As in SharedSparkSession: what CometDriverPlugin would install. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } + + conf diff --git a/dev/diffs/4.2.0.diff b/dev/diffs/4.2.0.diff index 01a29c48ae8..f58f210fbef 100644 --- a/dev/diffs/4.2.0.diff +++ b/dev/diffs/4.2.0.diff @@ -78,7 +78,7 @@ index 46558134f41..7fa990be776 100644 org.apache.datasketches diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala -index 7c5e49f3a69..28f83381fa0 100644 +index 7c5e49f3a69..1091462f8ef 100644 --- a/project/SparkBuild.scala +++ b/project/SparkBuild.scala @@ -455,9 +455,11 @@ object SparkBuild extends PomBuild { @@ -93,7 +93,7 @@ index 7c5e49f3a69..28f83381fa0 100644 enable(HiveThriftServer.settings)(hiveThriftServer) -@@ -1426,6 +1428,25 @@ object SQL { +@@ -1426,6 +1428,35 @@ object SQL { } } @@ -101,16 +101,26 @@ index 7c5e49f3a69..28f83381fa0 100644 + * Comet disables itself when its shuffle manager is not registered. When Comet is enabled, make + * that shuffle manager the default for every test SparkConf in the modules that have Comet on + * their classpath, so that suites which build their own SparkSession or SparkContext run Comet -+ * too, not just the ones that go through SharedSparkSession or TestHive. ++ * too, not just the ones that go through SharedSparkSession or TestHive. Comet's cache serializer ++ * is made the default the same way. Spark keeps the first cache serializer it loads for the rest ++ * of the JVM, so a suite that caches through a session of its own would otherwise leave the ++ * suites after it in Spark's cache format. + */ +object CometTestSettings { + private val cometEnabled = + sys.env.get("ENABLE_COMET").forall(v => v == "1" || v.equalsIgnoreCase("true")) + ++ private val defaults = Seq( ++ "spark.shuffle.manager" -> ++ "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager", ++ "spark.sql.cache.serializer" -> ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") ++ + lazy val settings: Seq[Setting[_]] = -+ if (cometEnabled && !sys.props.contains("spark.shuffle.manager")) { -+ Seq((Test / javaOptions) += -+ "-Dspark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") ++ if (cometEnabled) { ++ defaults.filterNot { case (key, _) => sys.props.contains(key) }.map { ++ case (key, value) => (Test / javaOptions) += s"-D$key=$value" ++ } + } else { + Seq.empty + } @@ -418,11 +428,41 @@ index 21a3ce1e122..f4762ab98f0 100644 SET spark.sql.ansi.enabled = false; -- In COMPENSATION views get invalidated if the type can't cast +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CacheTableInKryoSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CacheTableInKryoSuite.scala +index 72a2da16054..54146036f8a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/CacheTableInKryoSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/CacheTableInKryoSuite.scala +@@ -30,9 +30,16 @@ import org.apache.spark.storage.StorageLevel + class CacheTableInKryoSuite extends SharedSparkSession { + + override def sparkConf: SparkConf = { +- super.sparkConf ++ val conf = super.sparkConf + .set("spark.kryo.registrationRequired", "true") + .set("spark.serializer", "org.apache.spark.serializer.KryoSerializer") ++ // Comet's cache format needs its classes registered too, which is what Comet asks of any ++ // application that sets registrationRequired. ++ if (isCometEnabled) { ++ conf.set("spark.kryo.registrator", "org.apache.comet.CometKryoRegistrator") ++ } else { ++ conf ++ } + } + + test("SPARK-51777: sql.columnar.* classes registered in KryoSerializer") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala -index 085dbcd8046..3090d321b6c 100644 +index 085dbcd8046..2569c53815f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala -@@ -49,7 +49,7 @@ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEProp +@@ -37,6 +37,7 @@ import org.apache.spark.sql.catalyst.analysis.TempTableAlreadyExistsException + import org.apache.spark.sql.catalyst.expressions.SubqueryExpression + import org.apache.spark.sql.catalyst.plans.logical.{BROADCAST, Join, JoinStrategyHint, SHUFFLE_HASH} + import org.apache.spark.sql.catalyst.util.DateTimeConstants ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.connector.catalog.BasicInMemoryTableCatalog + import org.apache.spark.sql.connector.catalog.CatalogPlugin + import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.CatalogHelper +@@ -49,7 +50,7 @@ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEProp import org.apache.spark.sql.execution.columnar._ import org.apache.spark.sql.execution.command.CommandUtils import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation @@ -431,7 +471,28 @@ index 085dbcd8046..3090d321b6c 100644 import org.apache.spark.sql.execution.ui.SparkListenerSQLAdaptiveExecutionUpdate import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -@@ -564,7 +564,8 @@ class CachedTableSuite extends SharedSparkSession +@@ -127,6 +128,10 @@ class CachedTableSuite extends SharedSparkSession + case inMemoryTable @ InMemoryTableScanExec(_, _, relation) => + getNumInMemoryTablesRecursively(relation.cachedPlan) + + getNumInMemoryTablesInSubquery(inMemoryTable) + 1 ++ // Comet's cache scan keeps the predicates pushed into it on its original plan, out of reach ++ // of subquery reuse. The filter above it evaluates the same subqueries, and counts them. ++ case cometScan: CometInMemoryTableScanExec => ++ getNumInMemoryTablesRecursively(cometScan.originalPlan.relation.cachedPlan) + 1 + case p => + getNumInMemoryTablesInSubquery(p) + }.sum +@@ -407,7 +412,8 @@ class CachedTableSuite extends SharedSparkSession + assert(isExpectStorageLevel(rddId, Disk)) + } + +- test("InMemoryRelation statistics") { ++ test("InMemoryRelation statistics", ++ IgnoreComet("Comet's cache format reports Arrow buffer sizes")) { + sql("CACHE TABLE testData") + spark.table("testData").queryExecution.withCachedData.collect { + case cached: InMemoryRelation => +@@ -564,7 +570,8 @@ class CachedTableSuite extends SharedSparkSession df.collect() } assert( @@ -441,7 +502,19 @@ index 085dbcd8046..3090d321b6c 100644 } test("A cached table preserves the partitioning and ordering of its cached SparkPlan") { -@@ -1703,9 +1704,18 @@ class CachedTableSuite extends SharedSparkSession +@@ -1632,7 +1639,10 @@ class CachedTableSuite extends SharedSparkSession + checkAnswer(spark.table(tableName), Row(LocalDateTime.parse("2021-01-01T00:00:00"))) + spark.table(tableName).queryExecution.withCachedData.collect { + case cached: InMemoryRelation => +- assert(cached.stats.sizeInBytes === 8) ++ // Comet's cache format reports Arrow buffer sizes. ++ if (!isCometEnabled) { ++ assert(cached.stats.sizeInBytes === 8) ++ } + } + sql(s"UNCACHE TABLE $tableName") + } +@@ -1703,9 +1713,19 @@ class CachedTableSuite extends SharedSparkSession _.nodeName.contains("TableCacheQueryStage")) val aqeNode = findNodeInSparkPlanInfo(inMemoryScanNode.get, _.nodeName.contains("AdaptiveSparkPlan")) @@ -449,22 +522,47 @@ index 085dbcd8046..3090d321b6c 100644 - _.nodeName.contains("ResultQueryStage")) - aqePlanRoot.get.children.head.nodeName == "AQEShuffleRead" + // Spark 4.0 wraps results in ResultQueryStage. The coalescing indicator is AQEShuffleRead -+ // as the direct child of InputAdapter. ++ // as its direct child, as in Spark, when the table is cached in Comet's columnar format, ++ // or as the direct child of InputAdapter when Comet's plan is cached in Spark's format. + // AdaptiveSparkPlan -> ResultQueryStage -> WholestageCodegen -> + // CometColumnarToRow -> InputAdapter -> AQEShuffleRead (if coalesced) + val resultStage = aqeNode.get.children.head // ResultQueryStage -+ val wsc = resultStage.children.head // WholeStageCodegen ++ val wsc = resultStage.children.head // WholeStageCodegen or AQEShuffleRead + val c2r = wsc.children.head // ColumnarToRow or CometColumnarToRow + val inputAdapter = c2r.children.head // InputAdapter -+ resultStage.nodeName == "ResultQueryStage" && ++ resultStage.nodeName == "ResultQueryStage" && (wsc.nodeName == "AQEShuffleRead" || + wsc.nodeName.startsWith("WholeStageCodegen") && // could be "WholeStageCodegen (1)" + (c2r.nodeName == "CometColumnarToRow" || c2r.nodeName == "ColumnarToRow") && -+ inputAdapter.children.head.nodeName == "AQEShuffleRead" ++ inputAdapter.children.head.nodeName == "AQEShuffleRead") } withTempView("t0", "t1", "t2") { +@@ -1794,6 +1814,7 @@ class CachedTableSuite extends SharedSparkSession + val cached = spark.table("t") + val tableCache = collect(cached.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c.originalPlan + } + if (expected == StorageLevel.NONE) { + assert(tableCache.isEmpty) +@@ -2660,6 +2681,7 @@ class CachedTableSuite extends SharedSparkSession + + val inMemoryTableScan = collect(df.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c.originalPlan + } + assert(inMemoryTableScan.size == 1) + checkAnswer(df, Row(5) :: Nil) +@@ -2687,6 +2709,7 @@ class CachedTableSuite extends SharedSparkSession + + val subqueryInMemoryTableScan = collect(cteInSubquery.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c.originalPlan + } + assert(subqueryInMemoryTableScan.size == 1) + checkAnswer(cteInSubquery, Row(1) :: Nil) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala -index 5b8154d2900..eb3de47f387 100644 +index 5b8154d2900..d42374b127f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -30,8 +30,9 @@ import org.apache.spark.sql.catalyst.util.AUTO_GENERATED_ALIAS @@ -478,18 +576,7 @@ index 5b8154d2900..eb3de47f387 100644 import org.apache.spark.sql.expressions.Window import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -@@ -667,7 +668,9 @@ class DataFrameAggregateSuite extends SharedSparkSession - df.selectExpr("sort_array(collect_set(b) RESPECT NULLS)"), Seq(Row(Seq(null, 2)))) - } - -- test("SPARK-57298: collect_set normalizes NaN and -0.0 for floating-point types") { -+ test("SPARK-57298: collect_set normalizes NaN and -0.0 for floating-point types", -+ // Comet's native collect_set does not yet normalize NaN / -0.0. -+ IgnoreComet("https://github.com/apache/datafusion-comet/issues/4966")) { - checkAnswer( - sql("SELECT collect_set(v) FROM VALUES (double('NaN')), (double('NaN')) AS t(v)"), - Row(Seq(Double.NaN))) -@@ -932,7 +935,9 @@ class DataFrameAggregateSuite extends SharedSparkSession +@@ -932,7 +933,9 @@ class DataFrameAggregateSuite extends SharedSparkSession case _ => false }.isDefined) } else { @@ -500,7 +587,7 @@ index 5b8154d2900..eb3de47f387 100644 } // test case for ObjectHashAggregate and SortAggregate -@@ -941,7 +946,8 @@ class DataFrameAggregateSuite extends SharedSparkSession +@@ -941,7 +944,8 @@ class DataFrameAggregateSuite extends SharedSparkSession val objHashAggOrSortAggPlan = stripAQEPlan(objHashAggOrSortAggDF.queryExecution.executedPlan) if (useObjectHashAgg) { @@ -510,7 +597,7 @@ index 5b8154d2900..eb3de47f387 100644 } else { assert(objHashAggOrSortAggPlan.isInstanceOf[SortAggregateExec]) } -@@ -991,12 +997,12 @@ class DataFrameAggregateSuite extends SharedSparkSession +@@ -991,12 +995,12 @@ class DataFrameAggregateSuite extends SharedSparkSession assert(sortAggPlans.isEmpty) val objHashAggPlans = collect(aggPlan) { @@ -541,7 +628,7 @@ index 9733d51a91c..395a108abc8 100644 spark.range(100).write.saveAsTable(s"$dbName.$table2Name") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala -index d838ba4c234..cb0573d56d0 100644 +index d838ba4c234..95ba4076c4d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala @@ -23,10 +23,11 @@ import java.util.Locale @@ -557,7 +644,24 @@ index d838ba4c234..cb0573d56d0 100644 import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.{ExamplePoint, ExamplePointUDT, SharedSparkSession, SQLTestData} -@@ -1518,11 +1519,12 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP +@@ -1496,9 +1497,13 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP + + val union = df1.union(df2) + union.collect() +- checkIfColumnar(union.queryExecution.executedPlan, +- _.isInstanceOf[InMemoryTableScanExec], supported, 2) +- checkIfColumnar(union.queryExecution.executedPlan, _.isInstanceOf[UnionExec], supported, 1) ++ // Comet replaces the cache scans and the union with its own operators. ++ if (!isCometEnabled) { ++ checkIfColumnar(union.queryExecution.executedPlan, ++ _.isInstanceOf[InMemoryTableScanExec], supported, 2) ++ checkIfColumnar(union.queryExecution.executedPlan, ++ _.isInstanceOf[UnionExec], supported, 1) ++ } + checkAnswer(union, Row(1) :: Row(2) :: Row(3) :: Row(4) :: Row(5) :: Row(6) :: Nil) + + val nonColumnarUnion = df1.union(Seq(7, 8, 9).toDF("k")) +@@ -1518,11 +1523,12 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP val union = df1.repartition($"a").union(df2.repartition($"a")) val unionExec = union.queryExecution.executedPlan.collect { case u: UnionExec => u @@ -571,7 +675,7 @@ index d838ba4c234..cb0573d56d0 100644 } assert(shuffle.size == 1) -@@ -1553,11 +1555,12 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP +@@ -1553,11 +1559,12 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP val union = df1.repartition($"a").union(df2.repartition($"d")) val unionExec = union.queryExecution.executedPlan.collect { case u: UnionExec => u @@ -585,7 +689,7 @@ index d838ba4c234..cb0573d56d0 100644 } assert(shuffle.size == 1) -@@ -1572,10 +1575,10 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP +@@ -1572,10 +1579,10 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP // Avoid unnecessary shuffle if union output partitioning is enabled val shuffledUnion = union.repartition($"a") val shuffleNumBefore = union.queryExecution.executedPlan.collect { @@ -598,7 +702,7 @@ index d838ba4c234..cb0573d56d0 100644 } if (enabled) { -@@ -1604,6 +1607,7 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP +@@ -1604,6 +1611,7 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP val union = df1.repartitionByRange($"a").union(df2.repartitionByRange($"d")) val unionExec = union.queryExecution.executedPlan.collect { case u: UnionExec => u @@ -739,6 +843,99 @@ index f79824de8ff..5432984960f 100644 val df = Seq((1, "1"), (2, "2"), (1, "3"), (2, "4")).toDF("key", "value") val window = Window.partitionBy($"key").orderBy($"value") +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +index 0d1b0e1d981..e12c11e7b59 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetCacheSuite.scala +@@ -22,7 +22,8 @@ import java.time.LocalTime + import org.scalatest.concurrent.TimeLimits + import org.scalatest.time.SpanSugar._ + +-import org.apache.spark.sql.execution.ColumnarToRowExec ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec ++import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.columnar.{InMemoryRelation, InMemoryTableScanExec} + import org.apache.spark.sql.functions._ +@@ -38,6 +39,10 @@ class DatasetCacheSuite extends SharedSparkSession + with AdaptiveSparkPlanHelper { + import testImplicits._ + ++ // A scan of a cached relation, Spark's or Comet's. ++ private def isCacheScan(plan: SparkPlan): Boolean = ++ plan.isInstanceOf[InMemoryTableScanExec] || plan.isInstanceOf[CometInMemoryTableScanExec] ++ + /** + * Asserts that a cached [[Dataset]] will be built using the given number of other cached results. + */ +@@ -45,7 +50,7 @@ class DatasetCacheSuite extends SharedSparkSession + val plan = df.queryExecution.withCachedData + assert(plan.isInstanceOf[InMemoryRelation]) + val internalPlan = plan.asInstanceOf[InMemoryRelation].cacheBuilder.cachedPlan +- assert(find(internalPlan)(_.isInstanceOf[InMemoryTableScanExec]).size ++ assert(find(internalPlan)(isCacheScan).size + == numOfCachesDependedUpon) + } + +@@ -255,7 +260,7 @@ class DatasetCacheSuite extends SharedSparkSession + case i: InMemoryRelation => i.cacheBuilder.cachedPlan + } + assert(df2LimitInnerPlan.isDefined && +- !df2LimitInnerPlan.get.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ !df2LimitInnerPlan.get.exists(isCacheScan)) + } + + test("SPARK-27739 Save stats from optimized plan") { +@@ -288,14 +293,14 @@ class DatasetCacheSuite extends SharedSparkSession + val unionDf = df1.union(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.union(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("positive: union by name") { + val unionDf = df1.unionByName(df2).select($"i") + unionDf.cache() + val finalDf = unionDf.unionByName(df3.select($"i")) +- assert(finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by position") { +@@ -303,7 +308,7 @@ class DatasetCacheSuite extends SharedSparkSession + unionDf.cache() + val finalDf = unionDf.union(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + + withClue("negative: union by name") { +@@ -311,7 +316,7 @@ class DatasetCacheSuite extends SharedSparkSession + unionDf.cache() + val finalDf = unionDf.unionByName(df3) + // It's by design to break caching here. +- assert(!finalDf.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!finalDf.queryExecution.executedPlan.exists(isCacheScan)) + } + } + } +@@ -323,12 +328,12 @@ class DatasetCacheSuite extends SharedSparkSession + df1.cache() + // This is exactly the same as df1. + val df2 = spark.range(5).select(struct($"id".as("name", metadata))) +- assert(df2.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(df2.queryExecution.executedPlan.exists(isCacheScan)) + + val metadata2 = Metadata.fromJson("""{"k2": "v2"}""") + // Same with df1 except for the Alias metadata + val df3 = spark.range(5).select(struct($"id".as("name", metadata2))) +- assert(!df3.queryExecution.executedPlan.exists(_.isInstanceOf[InMemoryTableScanExec])) ++ assert(!df3.queryExecution.executedPlan.exists(isCacheScan)) + } + + test("SPARK-53418: Handle TimeType in ColumnAccessor") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala index 879569045b6..f3ff89067d2 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala @@ -762,10 +959,10 @@ index 879569045b6..f3ff89067d2 100644 assert(exchanges.size == 2) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -index c68f64d52f2..32489f2cb15 100644 +index c68f64d52f2..59f9f9bb092 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala -@@ -22,6 +22,7 @@ import org.scalatest.GivenWhenThen +@@ -22,9 +22,11 @@ import org.scalatest.GivenWhenThen import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression} import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode._ import org.apache.spark.sql.catalyst.plans.ExistenceJoin @@ -773,7 +970,11 @@ index c68f64d52f2..32489f2cb15 100644 import org.apache.spark.sql.connector.catalog.{InMemoryTableCatalog, InMemoryTableWithV2FilterCatalog} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive._ -@@ -192,6 +193,7 @@ abstract class DynamicPartitionPruningSuiteBase ++import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec + import org.apache.spark.sql.execution.datasources.v2.BatchScanExec + import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ReusedExchangeExec} + import org.apache.spark.sql.execution.joins.BroadcastHashJoinExec +@@ -192,6 +194,7 @@ abstract class DynamicPartitionPruningSuiteBase } val subqueryBroadcast = dpExprs.collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b @@ -781,7 +982,19 @@ index c68f64d52f2..32489f2cb15 100644 } val hasFilter = if (withSubquery) "Should" else "Shouldn't" -@@ -246,6 +248,8 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -228,7 +231,10 @@ abstract class DynamicPartitionPruningSuiteBase + } + + val isMainQueryAdaptive = plan.isInstanceOf[AdaptiveSparkPlanExec] +- subqueriesAll(plan).filterNot(subqueryBroadcast.contains).foreach { s => ++ // Comet's cache scan exposes Spark's own scan as its one subquery, for Spark's UI to draw ++ // the cached plan below it, but the query never runs it as a subquery. ++ subqueriesAll(plan).filterNot(subqueryBroadcast.contains) ++ .filterNot(_.isInstanceOf[InMemoryTableScanExec]).foreach { s => + val subquery = s match { + case r: ReusedSubqueryExec => r.child + case o => o +@@ -246,6 +252,8 @@ abstract class DynamicPartitionPruningSuiteBase val buf = collectDynamicPruningExpressions(df.queryExecution.executedPlan).collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b.indices.map(idx => b.buildKeys(idx)) @@ -790,7 +1003,7 @@ index c68f64d52f2..32489f2cb15 100644 } assert(buf.distinct.size == n) } -@@ -261,6 +265,12 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -261,6 +269,12 @@ abstract class DynamicPartitionPruningSuiteBase case s: BatchScanExec => s.runtimeFilters.collect { case d: DynamicPruningExpression => d.child } @@ -803,7 +1016,7 @@ index c68f64d52f2..32489f2cb15 100644 case _ => Nil } } -@@ -1203,10 +1213,16 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1203,10 +1217,16 @@ abstract class DynamicPartitionPruningSuiteBase val plan = df.queryExecution.executedPlan val countSubqueryBroadcasts = @@ -822,7 +1035,7 @@ index c68f64d52f2..32489f2cb15 100644 assert(countSubqueryBroadcasts == 1) assert(countReusedSubqueryBroadcasts == 1) -@@ -1578,6 +1594,7 @@ abstract class DynamicPartitionPruningSuiteBase +@@ -1578,6 +1598,7 @@ abstract class DynamicPartitionPruningSuiteBase val subqueryBroadcastExecs = collectWithSubqueries(df.queryExecution.executedPlan) { case s: SubqueryBroadcastExec => s @@ -830,7 +1043,7 @@ index c68f64d52f2..32489f2cb15 100644 } assert(subqueryBroadcastExecs.size === 1) subqueryBroadcastExecs.foreach { subqueryBroadcastExec => -@@ -1730,6 +1747,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat +@@ -1730,6 +1751,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat case s: BatchScanExec => // we use f1 col for v2 tables due to schema pruning s.output.exists(_.exists(_.argString(maxFields = 100).contains("f1"))) @@ -842,7 +1055,7 @@ index c68f64d52f2..32489f2cb15 100644 } assert(scanOption.isDefined) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala -index 618f0ab675e..be186f99b66 100644 +index 618f0ab675e..efcd4853251 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala @@ -268,7 +268,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite @@ -855,6 +1068,15 @@ index 618f0ab675e..be186f99b66 100644 withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true") { withTempView("df") { val df1 = spark.range(1, 100) +@@ -339,7 +340,7 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite + + val sqlText = "EXPLAIN FORMATTED SELECT (SELECT min(id) FROM df) as v" + val expected_pattern = +- "Subquery:1 Hosting operator id = 2 Hosting Expression = Subquery subquery#x" ++ "Subquery:1 Hosting operator id = \\d+ Hosting Expression = Subquery subquery#x" + + withNormalizedExplain(sqlText) { normalizedOutput => + assert(expected_pattern.r.findAllMatchIn(normalizedOutput).length == 1) @@ -471,7 +472,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite } } @@ -1373,10 +1595,18 @@ index 291aa7cab72..7783c37683e 100644 super.test(testName, testTags: _*) { withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala -index da6f6aca2ad..c02b7c99490 100644 +index da6f6aca2ad..64c981ad3d7 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala -@@ -1529,7 +1529,8 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper +@@ -40,6 +40,7 @@ import org.apache.spark.sql.catalyst.parser.ParseException + import org.apache.spark.sql.catalyst.plans.logical.{LocalLimit, Project, RepartitionByExpression, Sort} + import org.apache.spark.sql.connector.catalog.CatalogManager + import org.apache.spark.sql.connector.catalog.CatalogManager.SESSION_CATALOG_NAME ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.{CommandResultExec, OneRowRelationExec, UnionExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.execution.aggregate._ +@@ -1529,7 +1530,8 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper checkAnswer(sql("select -0.001"), Row(BigDecimal("-0.001"))) } @@ -1386,7 +1616,7 @@ index da6f6aca2ad..c02b7c99490 100644 AccumulatorSuite.verifyPeakExecutionMemorySet(sparkContext, "external sort") { sql("SELECT * FROM testData2 ORDER BY a ASC, b ASC").collect() } -@@ -1990,8 +1991,15 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper +@@ -1990,8 +1992,15 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper countAcc.add(1) x }) @@ -1403,6 +1633,24 @@ index da6f6aca2ad..c02b7c99490 100644 verifyCallCount( df.selectExpr("testUdf(a + 1) + testUdf(1 + a)", "testUdf(a + 1)"), Row(4, 2), 1) +@@ -3283,6 +3292,7 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper + + val inMemoryTableScan = collect(queryDf.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case c: CometInMemoryTableScanExec => c + } + assert(inMemoryTableScan.size == 2) + checkAnswer(queryDf, Row(0, 1) :: Row(1, 2) :: Row(2, 3) :: Row(3, 4) :: Row(4, 5) :: Nil) +@@ -4755,7 +4765,8 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper + ) + } + +- test("SPARK-47939: Explain should work with parameterized queries") { ++ test("SPARK-47939: Explain should work with parameterized queries", ++ IgnoreComet("Comet explain output is different")) { + def checkQueryPlan(df: DataFrame, plan: String): Unit = assert( + df.collect() + .map(_.getString(0)) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala index 395cb67f441..33ac6ed19af 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala @@ -1469,7 +1717,7 @@ index 395cb67f441..33ac6ed19af 100644 } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala -index bfcf583a705..6c177bff0e3 100644 +index bfcf583a705..749e6d7f1a4 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala @@ -217,7 +217,15 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp @@ -1488,7 +1736,15 @@ index bfcf583a705..6c177bff0e3 100644 val extensions = create { extensions => // inject rule that will run during AQE query stage preparation and will add custom tags // to the plan -@@ -311,6 +319,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp +@@ -260,6 +268,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp + df.select("i").filter($"i" > 1).cache() + assert(find(df.filter($"i" > 1).select("i").queryExecution.executedPlan) { + case _: org.apache.spark.sql.execution.columnar.InMemoryTableScanExec => true ++ case _: org.apache.spark.sql.comet.CometInMemoryTableScanExec => true + case _ => false + }.isDefined) + } +@@ -311,6 +320,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp } private def testInjectColumnar(enableAQE: Boolean): Unit = { @@ -1496,7 +1752,7 @@ index bfcf583a705..6c177bff0e3 100644 def collectPlanSteps(plan: SparkPlan): Seq[Int] = plan match { case a: AdaptiveSparkPlanExec => assert(a.toString.startsWith("AdaptiveSparkPlan isFinalPlan=true")) -@@ -366,6 +375,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp +@@ -366,6 +376,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp } test("reset column vectors") { @@ -1504,7 +1760,7 @@ index bfcf583a705..6c177bff0e3 100644 val session = SparkSession.builder() .master("local[1]") .config(COLUMN_BATCH_SIZE.key, 2) -@@ -553,6 +563,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp +@@ -553,6 +564,7 @@ class SparkSessionExtensionSuite extends PlanTest with SQLHelper with AdaptiveSp } test("SPARK-38697: Extend SparkSessionExtensions to inject rules into AQE Optimizer") { @@ -1527,6 +1783,33 @@ index d7b2511eac2..d5f5b940b94 100644 sc = new SparkContext("local[2]", "test") val session = classic.SparkSession.builder().sparkContext(sc).getOrCreate() import session.implicits._ +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +index 5222d5ce266..6f8f8bec3d4 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/StatisticsCollectionSuite.scala +@@ -847,14 +847,19 @@ class StatisticsCollectionSuite extends StatisticsCollectionTestBase with Shared + checkTableStats("t1", hasSizeInBytes = true, expectedRowCounts = None) + checkTableStats("t2", hasSizeInBytes = true, expectedRowCounts = None) + assert(getCatalogTable("v1").stats.isEmpty) +- checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ // Comet's cache format reports Arrow buffer sizes for a materialized cached view. ++ if (!isCometEnabled) { ++ checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ } + checkOptimizedPlanStats(spark.table("v2"), 1, None, Seq.empty) + + sql("ANALYZE TABLES COMPUTE STATISTICS") + checkTableStats("t1", hasSizeInBytes = true, expectedRowCounts = Some(10)) + checkTableStats("t2", hasSizeInBytes = true, expectedRowCounts = Some(20)) +- checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) +- checkOptimizedPlanStats(spark.table("v2"), 4, Some(1), Seq.empty) ++ if (!isCometEnabled) { ++ checkOptimizedPlanStats(spark.table("v1"), 4, Some(1), Seq.empty) ++ checkOptimizedPlanStats(spark.table("v2"), 4, Some(1), Seq.empty) ++ } + } + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala index 63589472854..f8c07a9b037 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala @@ -1726,7 +2009,7 @@ index 2d26356890d..2c5994f5fbc 100644 } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala -index 37684c7fce3..f3574dec867 100644 +index 37684c7fce3..a0a2f72b6c9 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala @@ -26,6 +26,8 @@ import org.apache.spark.sql.{AnalysisException, Row} @@ -1734,7 +2017,7 @@ index 37684c7fce3..f3574dec867 100644 import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.util.CollationFactory +import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashJoinExec, CometSortMergeJoinExec} -+import org.apache.spark.sql.comet.CometHashAggregateExec ++import org.apache.spark.sql.comet.{CometHashAggregateExec, CometInMemoryTableScanExec} import org.apache.spark.sql.connector.{DatasourceV2SQLBase, FakeV2ProviderWithCustomSchema} import org.apache.spark.sql.connector.catalog.{CatalogV2Util, Identifier, InMemoryTable} import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.CatalogHelper @@ -1793,6 +2076,15 @@ index 37684c7fce3..f3574dec867 100644 }.head.isInstanceOf[ArrayTransform]) } } +@@ -2005,7 +2017,7 @@ class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper { + // Checks in-memory fetching code path. + val all = sql("SELECT col FROM tbl") + assert(all.queryExecution.executedPlan.collectFirst { +- case _: InMemoryTableScanExec => true ++ case _: InMemoryTableScanExec | _: CometInMemoryTableScanExec => true + }.nonEmpty) + checkAnswer(all, Row("a")) + // Checks column stats code path. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala index 5ae23bc3338..5c2c3fff284 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala @@ -1961,7 +2253,7 @@ index ce549da03b4..47adf7cf94d 100644 // Fail to read ancient datetime values. withSQLConf(SQLConf.PARQUET_REBASE_MODE_IN_READ.key -> EXCEPTION.toString) { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala -index 3ef22ccb77e..54f67b45387 100644 +index 3ef22ccb77e..9aeb967617f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala @@ -56,6 +56,11 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper @@ -2021,23 +2313,6 @@ index 3ef22ccb77e..54f67b45387 100644 } } } -@@ -530,10 +541,12 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper - // Shuffle partition coalescing of the join is performed independent of the non-grouping - // aggregate on the other side of the union. - val finalPlan = stripAQEPlan(resultDf.queryExecution.executedPlan) -- assert( -- finalPlan.collect { -- case r @ CoalescedShuffleRead() => r -- }.size == 2) -+ if (!isCometEnabled) { -+ assert( -+ finalPlan.collect { -+ case r @ CoalescedShuffleRead() => r -+ }.size == 2) -+ } - } - withSparkSession(test, 100, None) - } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala index 60bf3a21e96..7693107f107 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/DataSourceScanExecRedactionSuite.scala @@ -2612,7 +2887,7 @@ index d70bd715879..074a9fa29d9 100644 }) checkAnswer(distinctWithId, Seq(Row(1, 0), Row(1, 0))) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala -index d6d19d21e65..85e2f441f81 100644 +index d6d19d21e65..7b3942d936b 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala @@ -27,13 +27,15 @@ import org.apache.spark.SparkException @@ -3088,7 +3363,18 @@ index d6d19d21e65..85e2f441f81 100644 }.isEmpty) assert(collect(initialExecutedPlan) { case i: InMemoryTableScanLike => i -@@ -3278,7 +3330,8 @@ class AdaptiveQueryExecSuite +@@ -3272,13 +3324,18 @@ class AdaptiveQueryExecSuite + + val (_, adaptive) = runAdaptiveAndVerifyResult( + "SELECT * FROM t1 WHERE c1 < (SELECT c2 FROM t2)") +- assert(findReusedSubquery(adaptive).size == 1) ++ // Comet's cache scan does not plan the subqueries in its pruning predicates, so only the ++ // filter above it runs this one and there is nothing to reuse. ++ if (!isCometEnabled) { ++ assert(findReusedSubquery(adaptive).size == 1) ++ } + } + } } } @@ -3098,7 +3384,7 @@ index d6d19d21e65..85e2f441f81 100644 val emptyDf = spark.range(1).where("false") val aggDf1 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df1")) val aggDf2 = emptyDf.agg(sum("id").as("id")).withColumn("name", lit("df2")) -@@ -3368,7 +3421,8 @@ class AdaptiveQueryExecSuite +@@ -3368,7 +3425,8 @@ class AdaptiveQueryExecSuite val plan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] assert(plan.inputPlan.isInstanceOf[TakeOrderedAndProjectExec]) @@ -3109,10 +3395,29 @@ index d6d19d21e65..85e2f441f81 100644 assert(o1.semanticEquals(o2), "Different output column order after AQE optimization") } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala -index 88be4adb6a4..f8fe831744e 100644 +index 88be4adb6a4..23ec5374f8a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/CachedBatchSerializerSuite.scala -@@ -228,9 +228,21 @@ class CachedBatchSerializerNoUnwrapSuite extends SharedSparkSession with Adaptiv +@@ -210,6 +210,18 @@ class CachedBatchSerializerNoUnwrapSuite extends SharedSparkSession with Adaptiv + classOf[DefaultCachedBatchSerializerNoUnwrap].getName) + } + ++ // Spark keeps the first cache serializer it loads for the rest of the JVM, which the suites ++ // before this one set to Comet's. Clear it on both sides, as CachedBatchSerializerSuite does. ++ protected override def beforeAll(): Unit = { ++ super.beforeAll() ++ clearSerializer() ++ } ++ ++ protected override def afterAll(): Unit = { ++ clearSerializer() ++ super.afterAll() ++ } ++ + test("Do not unwrap ColumnarToRowExec") { + withTempPath { workDir => + val workDirPath = workDir.getAbsolutePath +@@ -228,9 +240,21 @@ class CachedBatchSerializerNoUnwrapSuite extends SharedSparkSession with Adaptiv assert(cachedPlans.length == 2) cachedPlans.foreach { cachedPlan => @@ -3137,6 +3442,147 @@ index 88be4adb6a4..f8fe831744e 100644 } } } +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +index 57da12e8797..396d0745758 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +@@ -28,6 +28,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, + import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning + import org.apache.spark.sql.classic.DataFrame + import org.apache.spark.sql.columnar.CachedBatch ++import org.apache.spark.sql.comet.{CometFilterExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{FilterExec, InputAdapter, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions._ +@@ -509,11 +510,17 @@ class InMemoryColumnarQuerySuite extends SharedSparkSession with AdaptiveSparkPl + val planBeforeFilter = collect(df2.queryExecution.executedPlan) { + case f: FilterExec => f.child + case WholeStageCodegenExec(FilterExec(_, i: InputAdapter)) => i.child ++ case f: CometFilterExec => f.child + } +- assert(planBeforeFilter.head.isInstanceOf[InMemoryTableScanExec]) +- + val execPlan = planBeforeFilter.head +- assert(execPlan.executeCollectPublic().length == 0) ++ execPlan match { ++ // Comet's cache scan is columnar only, so count the rows of the batches it emits. ++ case c: CometInMemoryTableScanExec => ++ assert(c.executeColumnar().map(_.numRows().toLong).collect().sum == 0) ++ case _ => ++ assert(execPlan.isInstanceOf[InMemoryTableScanExec]) ++ assert(execPlan.executeCollectPublic().length == 0) ++ } + } + + test("SPARK-25727 - otherCopyArgs in InMemoryRelation does not include outputOrdering") { +@@ -543,7 +550,10 @@ class InMemoryColumnarQuerySuite extends SharedSparkSession with AdaptiveSparkPl + + // InMemoryRelation's stats is updated after materializing RDD + dfFromFile.collect() +- assert(inMemoryRelation.computeStats().sizeInBytes === 16) ++ // Comet's cache format reports Arrow buffer sizes. ++ if (!isCometEnabled) { ++ assert(inMemoryRelation.computeStats().sizeInBytes === 16) ++ } + + // test of catalog table + val dfFromTable = spark.catalog.createTable("table1", workDirPath).cache() +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +index 2c73622739a..5d0efeb263a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +@@ -18,6 +18,7 @@ + package org.apache.spark.sql.execution.columnar + + import org.apache.spark.SparkFunSuite ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.SparkPlan + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.functions.expr +@@ -37,11 +38,15 @@ class InMemoryRelationSuite extends SparkFunSuite + test("SPARK-47177: Cached SQL plan do not display final AQE plan in explain string") { + def findIMRInnerChild(p: SparkPlan): SparkPlan = { + val tableCache = find(p) { +- case _: InMemoryTableScanExec => true ++ case _: InMemoryTableScanExec | _: CometInMemoryTableScanExec => true + case _ => false + } + assert(tableCache.isDefined) +- tableCache.get.asInstanceOf[InMemoryTableScanExec].relation.innerChildren.head ++ val scan = tableCache.get match { ++ case c: CometInMemoryTableScanExec => c.originalPlan ++ case s => s.asInstanceOf[InMemoryTableScanExec] ++ } ++ scan.relation.innerChildren.head + } + + val d1 = spark.range(1).withColumn("key", expr("id % 100")) +diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +index 88ff51d0ff4..7d00fca3b0a 100644 +--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala ++++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala +@@ -17,6 +17,7 @@ + + package org.apache.spark.sql.execution.columnar + ++import org.apache.spark.sql.comet.CometInMemoryTableScanExec + import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + import org.apache.spark.sql.internal.SQLConf + import org.apache.spark.sql.test.SharedSparkSession +@@ -181,11 +182,16 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + val result = df.collect().map(_(0)).toArray + assert(result.length === 1) + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- assert(readPartitions === 5) +- assert(readBatches === 10) ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan has none of these test-only accumulators, so there is nothing to count. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ assert(scans.head.readPartitions.value === 5) ++ assert(scans.head.readBatches.value === 10) ++ } + } + + def checkBatchPruning( +@@ -202,14 +208,23 @@ class PartitionBatchPruningSuite extends SharedSparkSession with AdaptiveSparkPl + df.collect().map(_(0)).toArray + } + +- val (readPartitions, readBatches) = collect(df.queryExecution.executedPlan) { +- case in: InMemoryTableScanExec => (in.readPartitions.value, in.readBatches.value) +- }.head +- +- assert(readBatches === expectedReadBatches, s"Wrong number of read batches: $queryExecution") +- assert( +- readPartitions === expectedReadPartitions, +- s"Wrong number of read partitions: $queryExecution") ++ val scans = collect(df.queryExecution.executedPlan) { case in: InMemoryTableScanExec => in } ++ // Comet's cache scan prunes on the same statistics, but has none of these test-only ++ // accumulators to read, so only the answer above is checked for it. ++ if (scans.isEmpty) { ++ assert(collect(df.queryExecution.executedPlan) { ++ case c: CometInMemoryTableScanExec => c ++ }.nonEmpty) ++ } else { ++ val readPartitions = scans.head.readPartitions.value ++ val readBatches = scans.head.readBatches.value ++ assert( ++ readBatches === expectedReadBatches, ++ s"Wrong number of read batches: $queryExecution") ++ assert( ++ readPartitions === expectedReadPartitions, ++ s"Wrong number of read partitions: $queryExecution") ++ } + } + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala index fd8d1308e99..0e1f80045a3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala @@ -3623,10 +4069,10 @@ index b8f3ea3c6f3..bbd44221288 100644 val workDirPath = workDir.getAbsolutePath val input = spark.range(5).toDF("id") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -index 52746720eba..15df5e78440 100644 +index 52746720eba..98266bf06f5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala -@@ -21,7 +21,7 @@ import scala.reflect.ClassTag +@@ -21,13 +21,14 @@ import scala.reflect.ClassTag import org.apache.spark.AccumulatorSuite import org.apache.spark.internal.config.EXECUTOR_MEMORY @@ -3635,7 +4081,29 @@ index 52746720eba..15df5e78440 100644 import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BitwiseAnd, BitwiseOr, Cast, Expression, Literal, ShiftLeft} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} import org.apache.spark.sql.catalyst.plans.Inner -@@ -487,7 +487,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest + import org.apache.spark.sql.catalyst.plans.logical.BROADCAST + import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, PartitioningCollection} + import org.apache.spark.sql.classic.{Dataset, SparkSession} ++import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec} + import org.apache.spark.sql.execution.{DummySparkPlan, SparkPlan, WholeStageCodegenExec} + import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, DisableAdaptiveExecutionSuite, EnableAdaptiveExecutionSuite} + import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec +@@ -157,12 +158,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) +@@ -487,7 +490,8 @@ abstract class BroadcastJoinSuiteBase extends QueryTest } } @@ -3645,6 +4113,21 @@ index 52746720eba..15df5e78440 100644 withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "500") { val t1 = (0 until 100).map(i => (i % 5, i % 13)).toDF("i1", "j1") val t2 = (0 until 100).map(i => (i % 5, i % 14)).toDF("i2", "j2") +@@ -654,12 +658,14 @@ abstract class BroadcastJoinSuiteBase extends QueryTest + val df3 = df1.join(df2, Seq("key"), "inner") + val numCachedPlan = collect(df3.queryExecution.executedPlan) { + case i: InMemoryTableScanExec => i ++ case i: CometInMemoryTableScanExec => i + }.size + // df2 should be cached. + assert(numCachedPlan === 1) + + val numBroadCastHashJoin = collect(df3.queryExecution.executedPlan) { + case b: BroadcastHashJoinExec => b ++ case b: CometBroadcastHashJoinExec => b + }.size + // df2 should not be broadcasted. + assert(numBroadCastHashJoin === 0) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala index 2ae2ea6339f..30c8959a057 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala @@ -4518,10 +5001,10 @@ index e2c74533e7f..a12d55848ea 100644 val tblTargetName = "tbl_target" val tblSourceQualified = s"default.$tblSourceName" diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -index fb26d3311eb..13dc8b89a91 100644 +index fb26d3311eb..188b4d439a2 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala -@@ -102,6 +102,20 @@ trait SharedSparkSessionBase +@@ -102,6 +102,25 @@ trait SharedSparkSessionBase // this rule may potentially block testing of other optimization rules such as // ConstantPropagation etc. .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName) @@ -4538,6 +5021,11 @@ index fb26d3311eb..13dc8b89a91 100644 + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") + ++ // CometDriverPlugin installs Comet's cache serializer when ++ // spark.comet.exec.inMemoryCache.enabled is on, as it is by default. These sessions do not ++ // load the plugin, so install it here. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } conf.set( StaticSQLConf.WAREHOUSE_PATH, @@ -4616,10 +5104,10 @@ index 59022deaed7..f9aeacb5a9b 100644 override def beforeEach(): Unit = { super.beforeEach() diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala -index 8e7ff526a95..ea1072a7195 100644 +index 8e7ff526a95..518e1bae289 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/test/TestHive.scala -@@ -54,24 +54,34 @@ object TestHive +@@ -54,24 +54,38 @@ object TestHive new SparkContext( System.getProperty("spark.sql.test.master", "local[1]"), "TestSQLContext", @@ -4664,6 +5152,10 @@ index 8e7ff526a95..ea1072a7195 100644 + .set("spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.comet.shuffle.enabled", "true") ++ ++ // As in SharedSparkSession: what CometDriverPlugin would install. ++ conf.set("spark.sql.cache.serializer", ++ "org.apache.spark.sql.comet.execution.arrow.ArrowCachedBatchSerializer") + } + + conf diff --git a/docs/build.sh b/docs/build.sh index f4d839d8565..46a2130797f 100755 --- a/docs/build.sh +++ b/docs/build.sh @@ -31,6 +31,7 @@ rm -rf comet-0.15 rm -rf comet-0.16 rm -rf comet-0.17 rm -rf comet-1.0 +rm -rf comet-1.1 python3 generate-versions.py # Generate dynamic content (configs, compatibility matrices) for latest docs diff --git a/docs/generate-versions.py b/docs/generate-versions.py index 3ba6aee12e7..74a422ed375 100644 --- a/docs/generate-versions.py +++ b/docs/generate-versions.py @@ -154,6 +154,6 @@ def generate_docs(snapshot_version: str, latest_released_version: str, previous_ if __name__ == "__main__": print("Generating versioned user guide docs...") snapshot_version = get_version_from_pom() - latest_released_version = "1.0.0" - previous_versions = ["0.13.0", "0.14.0", "0.15.0", "0.16.0", "0.17.0"] + latest_released_version = "1.1.0" + previous_versions = ["0.13.0", "0.14.0", "0.15.0", "0.16.0", "0.17.0", "1.0.0"] generate_docs(snapshot_version, latest_released_version, previous_versions) \ No newline at end of file diff --git a/docs/source/about/faq.md b/docs/source/about/faq.md new file mode 100644 index 00000000000..25edfe2a816 --- /dev/null +++ b/docs/source/about/faq.md @@ -0,0 +1,209 @@ + + +# Frequently Asked Questions + +Answers to common questions about choosing, installing, and running Comet. If your question is not +answered here, see [Where can I ask questions or report bugs?](#where-can-i-ask-questions-or-report-bugs) + +## Choosing Comet + +### How does Comet compare to other open-source Spark accelerators? + +Several open-source projects speed up Spark by running query plans in native code. They differ in +the engine they build on and the hardware they need. + +- [Apache Gluten](https://gluten.apache.org/) is the closest to Comet in design: a Spark plugin that + passes the parts of a plan it supports to a native engine, such as the C++ engine Velox. See + [Comparison of Comet and Gluten](gluten_comparison.md). +- [Apache Auron (incubating)](https://auron.apache.org/), formerly Blaze, is also a Spark plugin + built on DataFusion. Comet is developed within the DataFusion project itself. +- The [NVIDIA cuDF plugin for Apache Spark](https://github.com/NVIDIA/cudf-spark), formerly the + RAPIDS Accelerator, runs queries on NVIDIA GPUs. Comet runs on the CPUs your cluster already has. + +Gluten and Auron are also adding Flink support, while Comet focuses on Spark alone. Which +accelerator is fastest depends on the workload, so we recommend benchmarking your own jobs. + +### Will Comet return the same results as Spark? + +That is the goal, and Comet runs Apache Spark's own SQL test suite in CI to check it. By default, +Comet uses implementations that match Spark, including running Spark's own generated code for +expressions whose native implementations behave differently. Those native implementations are +opt-in, one expression at a time. The known differences that remain, such as the handling of strings +that contain invalid UTF-8, are listed in the +[Compatibility Guide](../user-guide/latest/compatibility/index.md). If you find a difference that is +not documented, please [file an issue](https://github.com/apache/datafusion-comet/issues). + +### Which Spark and Java versions does Comet support? + +Spark 3.4, 3.5, 4.0, 4.1, and 4.2, on Java 17 or later. Spark 3.4 is deprecated. The +[installation guide](../user-guide/latest/installation.md#supported-spark-versions) lists the exact +Spark, Java, and Scala versions, and the +[Spark Version Compatibility](../user-guide/latest/compatibility/spark-versions.md) page lists known +issues for each Spark version. + +Comet drops a Spark version some time after the Spark project stops maintaining it. Spark 3.3 and +earlier are no longer supported, and Comet 1.1.0 dropped Java 11. Java 8 cannot be supported because +Apache Arrow's Java library no longer supports it. If you cannot upgrade Spark yet, stay on an earlier +Comet release that supports your version, as the +[versioning policy](versioning_policy.md#apache-spark-version-support) describes. + +### Does Comet work with vendor distributions of Spark? + +Comet is built and tested against open-source Apache Spark only. Spark distributions from cloud +providers and other vendors can change the internal Spark APIs that Comet calls and add operators +that Comet does not recognize, so Comet may fail with errors such as `NoSuchMethodError`, or leave +much of the work to Spark. To use Comet in the cloud, run open-source Apache Spark there, for example +on Kubernetes as described in the [Kubernetes guide](../user-guide/latest/kubernetes.md). + +### Which platforms do the published jars support? + +The jars in Maven Central include native libraries for Linux on amd64 and arm64, and need glibc 2.31 +or newer. They do not load on distributions with an older glibc, such as RHEL 7 and CentOS 7. On +those systems and on macOS, [build Comet from source](../user-guide/latest/source.md). + +For performance, the libraries target CPUs that are common in data centers: `x86-64-v3`, which +includes AVX2, on amd64, and Arm Neoverse N1 or newer cores on arm64. If Comet crashes with `SIGILL` +(illegal instruction), the CPU is older than that. Build Comet from source for it, and please +[open an issue](https://github.com/apache/datafusion-comet/issues) describing your environment. + +A jar you build yourself is smaller than the published one because it contains the native library +for one platform only. + +## Running Comet + +### Why do I get a `ClassNotFoundException` for `CometShuffleManager`? + +Spark could not find the Comet jar when it created the shuffle manager. On Spark 3.4 and 3.5, +executors create the shuffle manager before they load jars passed with `--jars` or `--packages`, so +the jar must be on the startup classpath. Either copy it into `$SPARK_HOME/jars` on every machine, or +set `spark.driver.extraClassPath` and `spark.executor.extraClassPath` to its path, as the +[installation guide](../user-guide/latest/installation.md#run-spark-shell-with-comet-enabled) shows. + +These two settings take a local file path, not an `hdfs://` or other URL. The JVM silently ignores +classpath entries that do not exist, so the jar must be at that path on every machine that runs a +driver or an executor. + +### How much memory does Comet need? + +Comet needs Spark's off-heap memory: without `spark.memory.offHeap.enabled=true`, Comet disables +itself and logs a warning. Comet's native operators share the off-heap pool, sized by +`spark.memory.offHeap.size`, with Spark. When the pool runs out, operators that can spill, such as +sorts, aggregations, and shuffle writes, spill to disk, and operators that cannot spill fail the +task. + +Comet also uses memory that the pool does not track, which has to fit in +`spark.executor.memoryOverhead`. If the cluster manager kills executors for exceeding their memory +limit, raise the overhead rather than the pool. Each executor logs its native memory use while Comet +runs, which you can use to size these settings. See +[Memory Tuning](../user-guide/latest/tuning/memory.md). + +### Why isn't my query faster with Comet? + +First check how much of the query Comet runs. Set `spark.comet.explain.fallback.enabled=true` to log +the reasons that parts of each query run in Spark, and see +[Understanding Comet Plans](../user-guide/latest/understanding-comet-plans.md) to read the plan. To +estimate how much of a workload Comet would accelerate without changing how it runs, set +`spark.comet.explain.planOnly.enabled=true`. + +Common causes of a small speedup are: + +- **Fallbacks inside a stage.** Each switch between Comet and Spark converts data between columnar + and row formats, which can cost more than Comet saves. See + [Reducing Row/Columnar Conversion Overhead](../user-guide/latest/tuning/transitions.md). +- **Comet's shuffle manager is not set.** If `spark.shuffle.manager` does not name Comet's shuffle + manager, Comet disables itself and logs a warning, unless `spark.comet.shuffle.enabled=false`, in + which case Comet runs but every shuffle runs in Spark. See + [Shuffle Tuning](../user-guide/latest/tuning/shuffle.md). +- **Too little memory.** Native operators spill to disk when the off-heap pool is too small. See + [Memory Tuning](../user-guide/latest/tuning/memory.md). +- **Little computation to accelerate.** Queries that spend most of their time reading from storage, + listing files, or planning, or that process little data, have less to gain. + +## Feature Support + +### Does Comet support Delta Lake, Apache Hudi, or Apache Paimon tables? + +Not yet. Comet accelerates [Apache Iceberg](../user-guide/latest/iceberg.md) tables, with a native +reader that is enabled by default and +[experimental native writes](../user-guide/latest/iceberg-writes.md). It does not accelerate scans of +Delta Lake, Hudi, or Paimon tables, so Spark reads those. Native Delta Lake reads are in development; +see the [roadmap](../contributor-guide/roadmap.md#delta-lake-support). Hudi and Paimon are not on the +roadmap. + +### Does Comet accelerate PySpark jobs and Python UDFs? + +PySpark DataFrame and SQL queries produce the same physical plans as Scala, so Comet accelerates them +the same way. Set Comet up as the [installation guide](../user-guide/latest/installation.md) +describes. There is no pip package. + +For Python UDFs, Comet has experimental support for `mapInArrow` and `mapInPandas` on Spark 4.0 and +later. It keeps their data in Arrow format instead of converting it to rows and back, and is disabled +by default. To enable it, set `spark.comet.exec.pyarrowUDF.enabled=true`; see +[PyArrow UDF Acceleration](../user-guide/latest/pyarrow-udfs.md). Scalar `@pandas_udf` functions are +not accelerated yet, and Python UDFs that do not use Arrow are not planned. + +### Does Comet plan to add an API for vectorized Java/Scala UDFs similar to the Rust UDF API? + +Not at the moment. Comet already runs ordinary Scala and Java UDFs in its pipeline without code +changes, compiling each call into a loop over a whole batch that the JVM optimizes well. See +[Scala UDF and Java UDF Support](../user-guide/latest/scala_java_udfs.md). + +A prototype ([#6697](https://github.com/apache/datafusion-comet/pull/6697)) found that, once the code +generator's overheads were reduced, a vectorized rewrite of a simple function such as `x + 1` saved +only about 3 ns per row. That does not justify asking users to rewrite their functions against +Comet's relocated Arrow classes and rebuild them for every Comet release, so the effort is going into +the code generator instead. + +If you have a use case that a function of one row handles poorly, such as working on the UTF-8 bytes +of strings or calling a library that processes whole batches, please describe it in +[#6694](https://github.com/apache/datafusion-comet/issues/6694). For native speed, see the +experimental [Rust UDF API](../user-guide/latest/rust_udfs.md). + +### Does Comet support Structured Streaming? + +No. Comet targets batch queries and leaves streaming queries to Spark, and streaming support is not +on the roadmap. See [Spark Operator Support](../user-guide/latest/operators.md#not-currently-planned). + +### Can I use Comet with a remote shuffle service? + +Yes, with [Apache Celeborn](https://celeborn.apache.org/). Comet accelerates scans, filters, and +other operators, and Celeborn's own Spark integration handles the shuffle. With the released Celeborn +0.6.x and 0.7.x clients, Comet's native shuffle cannot run through Celeborn. See +[Using Comet with Apache Celeborn](../user-guide/latest/celeborn.md). Native shuffle with Apache +Uniffle is not supported yet ([#4913](https://github.com/apache/datafusion-comet/issues/4913)). + +## Project and Community + +### Is there a roadmap? + +Yes. The [Comet Roadmap](../contributor-guide/roadmap.md) describes the major items that need +coordination between contributors. + +### Where can I ask questions or report bugs? + +- Ask questions in [GitHub Discussions](https://github.com/apache/datafusion-comet/discussions), or + in the Comet channels of the DataFusion + [Slack and Discord](https://datafusion.apache.org/contributor-guide/communication.html). +- Report bugs and request features in + [GitHub issues](https://github.com/apache/datafusion-comet/issues). +- Report security vulnerabilities privately, following the + [ASF security process](https://www.apache.org/security/), not in a public issue. +- Contributors also meet weekly on a video call; see + [Regular public meetings](../contributor-guide/contributing.md#regular-public-meetings). diff --git a/docs/source/about/gluten_comparison.md b/docs/source/about/gluten_comparison.md index 735662c52c4..56fb5340cc9 100644 --- a/docs/source/about/gluten_comparison.md +++ b/docs/source/about/gluten_comparison.md @@ -30,7 +30,9 @@ Comet and Gluten have very similar architectures. Both are Spark plugins that tr a serialized representation and pass the serialized plan to native code for execution. Gluten serializes the plans using the Substrait format and has an extensible architecture that supports execution -against multiple engines. Velox and Clickhouse are currently supported, but Velox is more widely used. +against multiple engines. Gluten 1.7.0 supports Velox and ClickHouse. A third backend, Bolt, a C++ engine from +ByteDance that is derived from Velox, has been added to Gluten's main branch but is not yet in a release. The rest of +this page compares Comet with Gluten's Velox backend. Comet serializes the plans in a proprietary Protocol Buffer format. Execution is delegated to Apache DataFusion. Comet does not plan to support multiple engines, but rather focus on a tight integration between Spark and DataFusion. @@ -87,7 +89,7 @@ Comet supports Spark 3.4, 3.5, 4.0, 4.1, and 4.2 in production builds. See the [Spark version compatibility guide]: /user-guide/latest/compatibility/spark-versions.md -Gluten supports Spark 3.3, 3.4, 3.5, 4.0, and 4.1. +Gluten 1.7.0 supports Spark 3.3, 3.4, 3.5, 4.0, and 4.1. ## ANSI Mode @@ -98,8 +100,10 @@ Comet implements ANSI semantics for the expressions it supports natively, includ ANSI cast behavior, and `try_*` variants. Queries running with `spark.sql.ansi.enabled=true` continue to be accelerated. See the [Comet Compatibility Guide] for details on which expressions have full ANSI coverage. -The Gluten Velox backend documents that ANSI mode is not supported and that any query executed with ANSI enabled -will fall back to vanilla Spark. See the [Gluten Velox limitations] page for the current status. +The Gluten Velox backend documents that ANSI mode is not supported: by default, any query executed with ANSI enabled +falls back to vanilla Spark. Setting `spark.gluten.sql.ansiFallback.enabled=false` makes Gluten attempt to run such +queries natively, but Gluten's developer documentation notes that the results do not yet match Spark. See the +[Gluten Velox limitations] page for the current status. [Gluten Velox limitations]: https://apache.github.io/gluten/velox-backend-limitations.html @@ -111,16 +115,22 @@ fraction of a workload that runs natively. Both projects can accelerate queries against Apache Iceberg tables, but they take different approaches and Gluten covers a broader set of table formats overall. -Comet provides a native Iceberg scan built on iceberg-rust. It has been tested with Iceberg 1.5 through 1.10 and -supports Iceberg spec v1 and v2, schema evolution, time travel and branch reads, positional and equality deletes -on merge-on-read tables, REST catalogs, and S3-compatible object storage. Iceberg writes still go through Spark. -Comet does not currently provide native integrations for Delta Lake, Hudi, or Paimon. See the -[Comet Iceberg guide] for the full list of supported features and known limitations. +Comet provides a native Iceberg scan built on iceberg-rust. It has been tested with Iceberg 1.5 and 1.8 through 1.11 +and supports Iceberg spec v1, v2, and v3, schema evolution, time travel and branch reads, positional and equality +deletes and deletion vectors, encrypted v3 tables, REST catalogs, and S3-compatible object storage. Comet can also +write Iceberg data files natively through iceberg-rust, as an experimental feature that is disabled by default (see +[Comet Iceberg writes]). Comet does not currently provide native integrations for Delta Lake, Hudi, or Paimon. See +the [Comet Iceberg guide] for the full list of supported features and known limitations. [Comet Iceberg guide]: /user-guide/latest/iceberg.md +[Comet Iceberg writes]: /user-guide/latest/iceberg-writes.md -Gluten ships dedicated modules for Iceberg, Delta Lake (2.0 through 4.0), Hudi, and Paimon. Users who need native -acceleration for Delta, Hudi, or Paimon will find broader coverage in Gluten today. +Gluten ships dedicated modules for Iceberg, Delta Lake (2.3 through 4.1, depending on the Spark version), Hudi, and +Paimon. Its [Iceberg documentation] says that reads of unpartitioned tables are offloaded, including tables with +position deletes, while reads of partitioned tables mostly fall back to Spark, as do equality deletes and all writes. +Users who need native acceleration for Delta, Hudi, or Paimon will find broader coverage in Gluten today. + +[Iceberg documentation]: https://github.com/apache/gluten/blob/v1.7.0/docs/get-started/VeloxIceberg.md ## Compatibility @@ -131,17 +141,17 @@ Spark are disabled by default, but users can opt in. See the [Comet Compatibilit [Comet Compatibility Guide]: https://datafusion.apache.org/comet/user-guide/latest/compatibility/index.html Gluten also aims to provide compatibility with Spark, and includes a subset of the Spark SQL tests in its own test -suite. See the Gluten [Velox backend limitations] page for known gaps, including notes on case sensitivity, regular -expression dialect (RE2 vs `java.util.regex`), NaN handling, and timestamp encodings. +suite. See the Gluten [Velox backend limitations] page for known gaps, such as differences in the regular expression +dialect (RE2 vs `java.util.regex`). [Velox backend limitations]: https://apache.github.io/gluten/velox-backend-limitations.html ## Codegen Dispatch -Comet has a feature called the codegen dispatcher that has no direct equivalent in Gluten. When an expression has a -native implementation with known semantic differences from Spark, or has no native implementation at all, Comet can -run Spark's own generated code for that expression inside the native pipeline. Data is passed to the JVM in Arrow -format, evaluated using Spark's byte-exact logic, and returned to the native pipeline for the rest of the query. +When an expression has a native implementation with known semantic differences from Spark, or has no native +implementation at all, Comet can run Spark's own generated code for that expression inside the native pipeline. Data +is passed to the JVM in Arrow format, evaluated using Spark's byte-exact logic, and returned to the native pipeline for +the rest of the query. Comet calls this the codegen dispatcher. This matters for two reasons: @@ -149,14 +159,15 @@ This matters for two reasons: expressions, JSON functions, and some datetime and array functions) route through the codegen dispatcher by default, so they produce results that match Spark exactly. The faster native path is opt-in per expression for users who accept its differences. -- **Fewer full fallbacks.** Because the dispatcher keeps evaluation inside the native pipeline instead of falling - back to whole-stage Spark execution, a single unsupported or incompatible expression does not force an entire - query stage back onto vanilla Spark. This keeps more of the plan running natively than a strict native-or-fallback - model would. +- **Fewer fallbacks.** Because the dispatcher keeps evaluation inside the native pipeline, a single unsupported or + incompatible scalar expression does not force its operator back onto vanilla Spark. -By contrast, Gluten falls back to vanilla Spark for expressions and operators that its Velox backend does not -support. See the [Comet Compatibility Guide] for more detail on how the codegen dispatcher works and which -expressions use it. +Gluten's closest equivalent is partial projection, which has been enabled by default since Gluten 1.3.0. When a +projection contains expressions that the backend cannot run, including Scala and Hive UDFs, Gluten splits the +projection so that vanilla Spark evaluates only those expressions and the rest of the projection runs natively. +Partial projection applies only to projections, so, for example, a filter that contains an unsupported expression +still falls back to Spark. See the [Comet Compatibility Guide] for more detail on how the codegen dispatcher works and +which expressions use it. ## Regular Expression Compatibility @@ -167,11 +178,13 @@ evaluate `rlike`, `regexp_replace`, `regexp_extract`, `regexp_extract_all`, `reg compatibility with Spark by default, including every pattern feature the JVM engine supports. Comet also offers a faster native Rust regex engine as an opt-in per expression for users who can accept its semantic differences. -Gluten's Velox backend evaluates regular expressions using the C++ RE2 engine. RE2 uses a different dialect from -`java.util.regex` and deliberately omits features such as backreferences and lookaround, so it cannot fully match -Spark's regex semantics regardless of configuration. This is a fundamental consequence of using a native C++ regex -engine rather than the JVM engine, and it is documented in the [Gluten Velox limitations] page. Comet's codegen -dispatcher sidesteps this class of incompatibility entirely by keeping Spark's own engine in the loop. +Gluten's Velox backend evaluates regular expressions using the C++ RE2 engine, which uses a different dialect from +`java.util.regex` and deliberately omits features such as backreferences and lookaround. Gluten falls back to Spark +for patterns that RE2 cannot compile, so those return correct results without acceleration. Patterns that RE2 accepts +run natively and can differ from Spark in edge cases: for example, `\s` does not match the vertical tab character. +Setting `spark.gluten.sql.fallbackRegexpExpressions=true` sends every regular expression function to Spark. See the +[Gluten Velox limitations] page for details. Comet's codegen dispatcher avoids this trade-off by running Spark's own +regex engine inside the native pipeline. ## Scala and Java UDF Support @@ -183,9 +196,10 @@ keep running natively while only the UDF itself is evaluated on the JVM. This co complex nested types and UDFs composed with other Catalyst expressions and higher-order functions. This means users can adopt Comet without giving up their existing Scala and Java UDFs, and without rewriting them in -the engine's native language. Gluten's Velox backend supports native C++ UDFs but does not run arbitrary existing -JVM UDFs in the native pipeline, so plans containing them typically fall back to Spark. See the -[Comet Scala and Java UDF guide] for the full list of supported and unsupported cases. +the engine's native language. Hive UDFs are not covered and still fall back to Spark. Gluten's Velox backend supports +native C++ UDFs, and its partial projection (see [Codegen Dispatch](#codegen-dispatch)) evaluates Scala and Hive UDFs +in a projection with vanilla Spark while the rest of the projection runs natively. A UDF in a filter makes the filter +fall back to Spark. See the [Comet Scala and Java UDF guide] for the full list of supported and unsupported cases. [Comet Scala and Java UDF guide]: /user-guide/latest/scala_java_udfs.md @@ -199,8 +213,15 @@ overall performance, with Gluten finishing roughly 9% faster than Comet across t The headline number masks wide per-query variation: some queries were significantly faster on Gluten (for example, large fact-table joins where Gluten uses a shuffled hash join strategy), while others were significantly faster on -Comet (for example, CPU-bound scan and aggregate pipelines). We expect Comet performance to continue improving over -time and for this gap to close. +Comet (for example, CPU-bound scan and aggregate pipelines). + +AWS Labs has also published separate runs of each accelerator against the same Spark baseline on the same TPC-DS 3TB +Parquet workload. [Comet 1.0.0] ran it 1.57× faster than Spark (and 1.75× faster on Iceberg tables), and +[Gluten 1.6.0] ran it 1.63× faster. We expect Comet performance to continue improving over time and for this gap to +close. + +[Comet 1.0.0]: https://awslabs.github.io/data-on-eks/docs/benchmarks/spark-datafusion-comet-benchmark +[Gluten 1.6.0]: https://awslabs.github.io/data-on-eks/docs/benchmarks/spark-gluten-velox-benchmark Although TPC-DS and TPC-H are good benchmarks for operators such as joins and aggregates, they don't necessarily represent real-world queries, especially for ETL use cases. For example, there are limited complex types involved @@ -216,7 +237,7 @@ management capabilities vs the complexities around installing C++ dependencies. Comet and Gluten are both good solutions for accelerating Spark jobs, and independent benchmarking shows they deliver similar overall performance. Comet currently has an edge for users on Spark 4.0 with ANSI mode enabled, for -users who want a fully native Iceberg scan path, and through its codegen dispatcher, which keeps partially compatible -or unsupported expressions running inside the native pipeline rather than falling back to Spark. Gluten holds a -small performance lead in the AWS Labs TPC-DS benchmark and offers broader native integration with Delta Lake, Hudi, -and Paimon, plus a second backend in ClickHouse. We recommend trying both to see which is the best fit for your needs. +Iceberg workloads, and through its codegen dispatcher, which runs partially compatible or unsupported expressions with +Spark's exact semantics inside the native pipeline. Gluten holds a small performance lead in the AWS Labs TPC-DS +benchmarks and offers broader native integration with Delta Lake, Hudi, and Paimon, plus a second backend in +ClickHouse. We recommend trying both to see which is the best fit for your needs. diff --git a/docs/source/about/index.md b/docs/source/about/index.md index c5bddb38d19..5573a9103c8 100644 --- a/docs/source/about/index.md +++ b/docs/source/about/index.md @@ -20,10 +20,14 @@ under the License. # About Background on the project: how Comet compares to similar accelerators, how releases -are versioned, and links to the Apache Software Foundation resources behind it. +are versioned, answers to common questions, and links to the Apache Software Foundation resources +behind it. -- [Comparison with Gluten](gluten_comparison.md) — how Comet differs from other Spark accelerators +- [Comparison with Gluten](gluten_comparison.md) — how Comet differs from Apache Gluten, the most + similar accelerator - [Versioning Policy](versioning_policy.md) — what our version numbers mean and our release cadence +- [Frequently Asked Questions](faq.md) — answers to common questions about choosing, installing, and + running Comet - [Blog Posts and Talks](blogs-and-talks.md) — external writing and conference talks about Comet - [ASF Links](../asf/index.md) — licensing, donations, security reporting, and the code of conduct @@ -34,6 +38,7 @@ are versioned, and links to the Apache Software Foundation resources behind it. Comparison with Gluten Versioning Policy +Frequently Asked Questions Blog Posts and Talks ASF Links <../asf/index> ``` diff --git a/docs/source/changelog/1.1.0.md b/docs/source/changelog/1.1.0.md new file mode 100644 index 00000000000..dfbaf5abda7 --- /dev/null +++ b/docs/source/changelog/1.1.0.md @@ -0,0 +1,503 @@ + + +# DataFusion Comet 1.1.0 Changelog + +This release consists of 416 commits from 41 contributors. See credits at the end of this changelog for more information. + +**Fixed bugs:** + +- fix: skip null slots when checking overflow in unary negation [#5162](https://github.com/apache/datafusion-comet/pull/5162) (Smallfu666) +- fix: surface next_day and make_date ANSI errors as Spark exceptions [#5167](https://github.com/apache/datafusion-comet/pull/5167) (peterxcli) +- fix: normalize nested field nullability in ShuffleScanExec and ExpandExec [#5138](https://github.com/apache/datafusion-comet/pull/5138) (andygrove) +- fix: propagate the Spark task ClassLoader to JVM UDF calls [#5282](https://github.com/apache/datafusion-comet/pull/5282) (andygrove) +- fix: avoid duplicate CheckOverflow evaluation for decimal division [#5225](https://github.com/apache/datafusion-comet/pull/5225) (peterxcli) +- fix: match Spark whitespace trimming in to_time and try_to_time [#5364](https://github.com/apache/datafusion-comet/pull/5364) (sunchao) +- fix: preserve Catalyst nullability and field IDs in native Parquet writes [#5369](https://github.com/apache/datafusion-comet/pull/5369) (sunchao) +- fix: guard against silent fail_on_error loss in scalar wiring (#5074) [#5359](https://github.com/apache/datafusion-comet/pull/5359) (sam-1112) +- fix: preserve Spark semantics for dictionary-encoded Parquet inputs and reject dictionary targets [#5234](https://github.com/apache/datafusion-comet/pull/5234) (peterxcli) +- fix: canonicalize NaN in flat arrays_overlap float keys [#5376](https://github.com/apache/datafusion-comet/pull/5376) (sunchao) +- fix: report native shuffle write metrics accurately [#5370](https://github.com/apache/datafusion-comet/pull/5370) (sunchao) +- fix: Native shuffle fails with a 2GB task serialization OOM on jobs with many partitions [#5392](https://github.com/apache/datafusion-comet/pull/5392) (parthchandra) +- fix: format NativeMemoryConsumer id in toString [#5398](https://github.com/apache/datafusion-comet/pull/5398) (ywskycn) +- fix: ignore reader-side parquet.hadoop.vectored.io.enabled in Iceberg native-write detection [#5410](https://github.com/apache/datafusion-comet/pull/5410) (snmvaughan) +- fix: support map casts with NullType elements [#5045](https://github.com/apache/datafusion-comet/pull/5045) (peterxcli) +- fix: distinguish reflection failure from absent accessor in IcebergReflection [#5412](https://github.com/apache/datafusion-comet/pull/5412) (unikdahal) +- fix: use current shuffle config in aggregate test [#5439](https://github.com/apache/datafusion-comet/pull/5439) (peterxcli) +- fix: support wide years in native make_date [#5443](https://github.com/apache/datafusion-comet/pull/5443) (peterxcli) +- fix: report native child spill metrics in shuffle tasks [#5445](https://github.com/apache/datafusion-comet/pull/5445) (sunchao) +- fix: expose native memory usage to Spark [#5408](https://github.com/apache/datafusion-comet/pull/5408) (ywskycn) +- fix: release native shuffle reservation after spill failure [#5461](https://github.com/apache/datafusion-comet/pull/5461) (peterxcli) +- fix: track peak memory before JVM shuffle spill [#5463](https://github.com/apache/datafusion-comet/pull/5463) (peterxcli) +- fix: RAII for memory pool registration [#5464](https://github.com/apache/datafusion-comet/pull/5464) (peterxcli) +- fix: tighten RSS writer visibility and JNI array limits [#5475](https://github.com/apache/datafusion-comet/pull/5475) (pingzh) +- fix: report native operator spill metrics in Spark task metrics for non-shuffle stages [#5497](https://github.com/apache/datafusion-comet/pull/5497) (peterxcli) +- fix: report bounded shuffle allocator memory usage [#5516](https://github.com/apache/datafusion-comet/pull/5516) (ywskycn) +- fix: use Spark type names in ANSI abs overflow errors [#5357](https://github.com/apache/datafusion-comet/pull/5357) (Smallfu666) +- fix: make collect_list/collect_set argument coercion a normalization barrier [#5159](https://github.com/apache/datafusion-comet/pull/5159) (andygrove) +- fix: make task-shared memory pool as ref-counted RAII guard [#5494](https://github.com/apache/datafusion-comet/pull/5494) (peterxcli) +- fix: accept UTC timezone aliases in Python Arrow input [#5556](https://github.com/apache/datafusion-comet/pull/5556) (sunchao) +- fix: normalize scalar float sort and window rank keys [#5469](https://github.com/apache/datafusion-comet/pull/5469) (sunchao) +- fix: make CometExplodeExec respect batch size [#5362](https://github.com/apache/datafusion-comet/pull/5362) (andygrove) +- fix: align Spark 4.2 Python worker configuration [#5561](https://github.com/apache/datafusion-comet/pull/5561) (sunchao) +- fix: prevent memory leak after failed Arrow vector import [#5539](https://github.com/apache/datafusion-comet/pull/5539) (1fanwang) +- fix: preserve Arrow Field metadata across C Data exports [#5552](https://github.com/apache/datafusion-comet/pull/5552) (peterxcli) +- fix: rebase map offsets in mapsort so sliced maps do not overrun entries [#5630](https://github.com/apache/datafusion-comet/pull/5630) (viirya) +- fix: scope Celeborn bootstrap hooks to Comet clients [#5627](https://github.com/apache/datafusion-comet/pull/5627) (pingzh) +- fix: skip codegen dispatcher null short-circuit when a foldable subtree can raise [#5623](https://github.com/apache/datafusion-comet/pull/5623) (andygrove) +- fix: accept dictionary encodings in remote shuffle [#5650](https://github.com/apache/datafusion-comet/pull/5650) (pingzh) +- fix: make CometDiskBlockWriter spill registry per-task instead of executor-global [#5493](https://github.com/apache/datafusion-comet/pull/5493) (peterxcli) +- fix: Bump iceberg-rust so native Iceberg writes URL-escape partition paths [#5651](https://github.com/apache/datafusion-comet/pull/5651) (andygrove) +- fix(celeborn): reject unsafe native push completion tracking [#5665](https://github.com/apache/datafusion-comet/pull/5665) (pingzh) +- fix: match Spark's null short-circuiting in array_join and enable it natively [#5558](https://github.com/apache/datafusion-comet/pull/5558) (Visorgood) +- fix: make native shuffle spill metrics independent of input batching [#5628](https://github.com/apache/datafusion-comet/pull/5628) (sunchao) +- fix: expand object store option references, uniquify constant metadata names, drop dead parquet JNI [#5653](https://github.com/apache/datafusion-comet/pull/5653) (dwsmith1983) +- fix: match Spark's ANSI bound check for float/double to integral casts [#5683](https://github.com/apache/datafusion-comet/pull/5683) (peterxcli) +- fix: return NULL from rpad/lpad when the length column is NULL instead of panicking [#5680](https://github.com/apache/datafusion-comet/pull/5680) (peterxcli) +- fix: fall back for concat_ws with array arguments instead of failing natively [#5679](https://github.com/apache/datafusion-comet/pull/5679) (peterxcli) +- fix: prevent silent overflow when reading Parquet TIMESTAMP_MILLIS values [#5177](https://github.com/apache/datafusion-comet/pull/5177) (peterxcli) +- fix: recover native Celeborn shuffle from oversized rows [#5668](https://github.com/apache/datafusion-comet/pull/5668) (pingzh) +- fix: report native shuffle read metrics [#5554](https://github.com/apache/datafusion-comet/pull/5554) (peterxcli) +- fix: correctly rounded decimal to double/float cast matching BigDecimal.doubleValue/floatValue [#5684](https://github.com/apache/datafusion-comet/pull/5684) (peterxcli) +- fix: restore columnar transitions under the native Iceberg write [#5696](https://github.com/apache/datafusion-comet/pull/5696) (andygrove) +- fix: make columnar-to-row benchmarks exercise Comet [#5718](https://github.com/apache/datafusion-comet/pull/5718) (rich7420) +- fix: distinguish "nothing spilled" from a spill backend with no local path [#5726](https://github.com/apache/datafusion-comet/pull/5726) (andygrove) +- fix: propagate Arrow array copy errors [#5747](https://github.com/apache/datafusion-comet/pull/5747) (rich7420) +- fix: keep the dictionary hash fast path off nested and reseeded buffers [#5757](https://github.com/apache/datafusion-comet/pull/5757) (viirya) +- fix: preserve ANSI errors for rejected TIMESTAMP_NTZ casts [#5752](https://github.com/apache/datafusion-comet/pull/5752) (peterxcli) +- fix: apply the parent struct's null mask before hashing its fields [#5754](https://github.com/apache/datafusion-comet/pull/5754) (viirya) +- fix: read Iceberg tables partitioned by an unknown transform [#5759](https://github.com/apache/datafusion-comet/pull/5759) (andygrove) +- fix: native Iceberg write panics on an evolved partition spec and on a timestamptz partition path [#5729](https://github.com/apache/datafusion-comet/pull/5729) (andygrove) +- fix: let AQE optimize queries over Comet caches [#5733](https://github.com/apache/datafusion-comet/pull/5733) (peterxcli) +- fix: dispatch Iceberg system functions wrapped as ApplyFunctionExpression [#5773](https://github.com/apache/datafusion-comet/pull/5773) (andygrove) +- fix: attach tokio runtime threads to the JVM as daemon threads [#5748](https://github.com/apache/datafusion-comet/pull/5748) (zhangfengcdt) +- fix: check nested TIMESTAMP_MILLIS overflow in unfiltered scans [#5740](https://github.com/apache/datafusion-comet/pull/5740) (peterxcli) +- fix: match iceberg-java's exception for unclustered input to a clustered Iceberg write [#5779](https://github.com/apache/datafusion-comet/pull/5779) (andygrove) +- fix: enable FIRST/LAST partial merge [#5041](https://github.com/apache/datafusion-comet/pull/5041) (peterxcli) +- fix: preserve aggregate result identity during exchange reuse [#5470](https://github.com/apache/datafusion-comet/pull/5470) (sunchao) +- fix: ignore structural tags when lifting expression coverage [#5471](https://github.com/apache/datafusion-comet/pull/5471) (sunchao) +- fix: align string to timestamp parsing with Spark's segment rules [#5682](https://github.com/apache/datafusion-comet/pull/5682) (peterxcli) +- fix: roll native Iceberg data files on iceberg-java's 1000-row grid [#5780](https://github.com/apache/datafusion-comet/pull/5780) (andygrove) +- fix: decide libhdfs routing from the scheme as written [#5825](https://github.com/apache/datafusion-comet/pull/5825) (comphead) +- fix: list a fanout Iceberg write's data files in a stable order [#5810](https://github.com/apache/datafusion-comet/pull/5810) (andygrove) +- fix: isolate object-store registration by backend and configuration [#5503](https://github.com/apache/datafusion-comet/pull/5503) (sunchao) +- fix: Delete completed tasks' data files when an Iceberg write job fails [#5663](https://github.com/apache/datafusion-comet/pull/5663) (andygrove) +- fix: apply Spark's Parquet conversion rules to nested struct/list/map fields [#5681](https://github.com/apache/datafusion-comet/pull/5681) (peterxcli) +- fix: preserve nulls for Boolean/Byte/Short/Integer columns in FuzzDataGenerator [#5855](https://github.com/apache/datafusion-comet/pull/5855) (Smallfu666) +- fix: Accept explicit positive years in timestamp casts [#5858](https://github.com/apache/datafusion-comet/pull/5858) (peterxcli) +- fix: decline structs with duplicate field names before they reach Java Arrow [#5866](https://github.com/apache/datafusion-comet/pull/5866) (dwsmith1983) +- fix: explain ObjectHashAggregate fallback when Comet shuffle is disabled [#5746](https://github.com/apache/datafusion-comet/pull/5746) (0lai0) +- fix: preserve join and generator semantics in plan identity [#5828](https://github.com/apache/datafusion-comet/pull/5828) (ErikBPF) +- fix: propagate Parquet field-name folding failures [#5845](https://github.com/apache/datafusion-comet/pull/5845) (sunchao) +- fix: decode dictionary input for PyArrow UDFs [#5560](https://github.com/apache/datafusion-comet/pull/5560) (sunchao) +- fix: size JVM shuffle pointer array growth from the array, not the data pages [#5907](https://github.com/apache/datafusion-comet/pull/5907) (andygrove) +- fix: render float and double Iceberg partition values like iceberg-java [#5840](https://github.com/apache/datafusion-comet/pull/5840) (andygrove) +- fix: align time parsing and native second extraction with Spark [#5738](https://github.com/apache/datafusion-comet/pull/5738) (peterxcli) +- fix: normalize floating-point values in native collect_set [#5166](https://github.com/apache/datafusion-comet/pull/5166) (peterxcli) +- fix: keep Iceberg complex null checks on native scans [#5732](https://github.com/apache/datafusion-comet/pull/5732) (ErikBPF) +- fix: decode invalid UTF-8 at the JVM to native FFI import boundary [#5310](https://github.com/apache/datafusion-comet/pull/5310) (manuzhang, andygrove) +- fix: fall back when a struct repeats a Parquet field id [#6004](https://github.com/apache/datafusion-comet/pull/6004) (comphead) +- fix: normalize signed zero in nested float array comparisons [#5235](https://github.com/apache/datafusion-comet/pull/5235) (divyankshah) +- fix: fall back to Spark for native Iceberg writes to gs:// through HadoopFileIO [#5935](https://github.com/apache/datafusion-comet/pull/5935) (zhangfengcdt) +- fix: gate the regr_r2 degenerate-case swap on the Spark patch release [#6042](https://github.com/apache/datafusion-comet/pull/6042) (dwsmith1983) +- fix: include ABFS container in object store cache key [#5053](https://github.com/apache/datafusion-comet/pull/5053) (peterxcli) +- fix: preserve Spark row index read errors [#6046](https://github.com/apache/datafusion-comet/pull/6046) (liupoyi-1031) +- fix: remove the ineffective spark.executor.memoryOverhead adjustment from the driver plugin [#6054](https://github.com/apache/datafusion-comet/pull/6054) (andygrove) +- fix: count each memory pool once in analyze_trace [#5991](https://github.com/apache/datafusion-comet/pull/5991) (andygrove) +- fix: revert unsafe partial aggregates after final fallback [#5421](https://github.com/apache/datafusion-comet/pull/5421) (sunchao) +- fix: restore the site's mermaid diagrams and make a dropped one fail CI [#6064](https://github.com/apache/datafusion-comet/pull/6064) (andygrove) +- fix: write cached batches to the schema width, not the batch width [#6090](https://github.com/apache/datafusion-comet/pull/6090) (andygrove) +- fix(iceberg): guard native Iceberg scan driver-metric double-post, add metrics docs and tests [#6085](https://github.com/apache/datafusion-comet/pull/6085) (parthchandra) +- fix: let decimal SUM recover from an intermediate overflow like Spark [#6041](https://github.com/apache/datafusion-comet/pull/6041) (dwsmith1983) +- fix: Nested floating-point IN membership does not match Spark for signed zero [#6073](https://github.com/apache/datafusion-comet/pull/6073) (mizulun) +- fix: let Comet memory pools overcommit on grow instead of panicking [#6128](https://github.com/apache/datafusion-comet/pull/6128) (andygrove) +- fix: correct two nightly test failures on Spark 3.4 and 4.2 [#6156](https://github.com/apache/datafusion-comet/pull/6156) (andygrove) +- fix: support null calendar interval literals [#5133](https://github.com/apache/datafusion-comet/pull/5133) (peterxcli) +- fix: wrap Iceberg split-write failures the way Spark does when abort fails [#6153](https://github.com/apache/datafusion-comet/pull/6153) (andygrove) +- fix: support Utf8/LargeUtf8/Utf8View in native RLike without panicking [#5215](https://github.com/apache/datafusion-comet/pull/5215) (sam-1112) +- fix: normalize noncanonical NaN literals in comparisons [#5472](https://github.com/apache/datafusion-comet/pull/5472) (sunchao) +- fix: preserve map field metadata and honor target sorted flag in cast_map_to_map [#5227](https://github.com/apache/datafusion-comet/pull/5227) (Smallfu666) +- fix: reject duplicate Parquet field names before decoding [#5786](https://github.com/apache/datafusion-comet/pull/5786) (ErikBPF) +- fix: preserve Parquet conversion errors during join filtering [#6067](https://github.com/apache/datafusion-comet/pull/6067) (pingzh) +- fix: plan Iceberg writes with Spark's own operator when Comet is disabled [#6151](https://github.com/apache/datafusion-comet/pull/6151) (andygrove) +- fix: read shuffle write buffer, spill limit and off-heap sizes in bytes [#6191](https://github.com/apache/datafusion-comet/pull/6191) (andygrove) +- fix: bound shuffle schema cache retention and preserve eviction order [#6098](https://github.com/apache/datafusion-comet/pull/6098) (sunchao) +- fix: do not run Comet in on-heap mode without `spark.comet.exec.onHeap.enabled` [#6195](https://github.com/apache/datafusion-comet/pull/6195) (andygrove) +- fix: remove misleading native opt-in for dispatch-only datetime expressions [#6182](https://github.com/apache/datafusion-comet/pull/6182) (LinSimon-901101) +- fix: support CalendarIntervalType hashing [#5135](https://github.com/apache/datafusion-comet/pull/5135) (peterxcli) +- fix: report `Input` column when native Iceberg scan is enabled or native shuffle is enabled [#5265](https://github.com/apache/datafusion-comet/pull/5265) (hsiang-c) +- fix: skip the executor memory overhead warning when a factor is set or in local mode [#6198](https://github.com/apache/datafusion-comet/pull/6198) (andygrove) +- fix: keep operators above a cached relation native after the AQE re-plan [#6208](https://github.com/apache/datafusion-comet/pull/6208) (andygrove) +- fix: preserve current AQE logical-stage links on Comet operators [#5483](https://github.com/apache/datafusion-comet/pull/5483) (sunchao) +- fix: check each fair_unified reservation against its own share [#6205](https://github.com/apache/datafusion-comet/pull/6205) (andygrove) +- fix(iceberg): don't push transform residuals as their source column, fail on residual errors [#6154](https://github.com/apache/datafusion-comet/pull/6154) (andygrove) +- fix: match Spark's duplicate field and field id semantics in parquet field lookup [#5654](https://github.com/apache/datafusion-comet/pull/5654) (dwsmith1983) +- fix: park the native scan loop instead of busy-polling while waiting on native I/O [#6219](https://github.com/apache/datafusion-comet/pull/6219) (andygrove, mixermt) +- fix: refresh S3 policy locations when a location's credential fails [#6223](https://github.com/apache/datafusion-comet/pull/6223) (andygrove, snmvaughan) +- fix: inject Comet session extension rules only once per session [#6230](https://github.com/apache/datafusion-comet/pull/6230) (andygrove) +- fix: [branch-1.1] reject a file without field ids at any depth whether or not id matching is on (#6116) [#6266](https://github.com/apache/datafusion-comet/pull/6266) (andygrove, dwsmith1983) +- fix: [branch-1.1] decline native Iceberg writes with a custom location provider (#6216) [#6305](https://github.com/apache/datafusion-comet/pull/6305) (andygrove, liupoyi-1031) +- fix: [branch-1.1] Native S3 scan on EKS/IRSA turns a transient STS throttle into a hard 403 storm (#6025) [#6323](https://github.com/apache/datafusion-comet/pull/6323) (andygrove, parthchandra) +- fix: [branch-1.1] log partial memory grants at DEBUG and drop the memory usage dump (#6269) [#6346](https://github.com/apache/datafusion-comet/pull/6346) (andygrove) +- fix: [branch-1.1] zero sliced boolean offsets at every level before exporting to the JVM (#6339) [#6449](https://github.com/apache/datafusion-comet/pull/6449) (andygrove) +- fix: [branch-1.1] build with Rust 1.99, which deprecates the legacy f64 constants and fetch_update (#6507) [#6526](https://github.com/apache/datafusion-comet/pull/6526) (andygrove) +- fix: [branch-1.1] make adaptive aggregation skipping opt-in (#6474) [#6488](https://github.com/apache/datafusion-comet/pull/6488) (andygrove) +- fix: [branch-1.1] fall back to Spark for \_metadata.file_block_start and file_block_length (#6510) [#6511](https://github.com/apache/datafusion-comet/pull/6511) (andygrove) +- fix: [branch-1.1] fall back for rank limits over nested float keys (#6468) [#6487](https://github.com/apache/datafusion-comet/pull/6487) (andygrove) +- fix: [branch-1.1] fall back for incompatible regression aggregates (#6451) [#6489](https://github.com/apache/datafusion-comet/pull/6489) (andygrove) +- fix: [branch-1.1] coerce native IF branches to a common type (#6458) [#6491](https://github.com/apache/datafusion-comet/pull/6491) (andygrove) +- fix: [branch-1.1] match pre-epoch Iceberg temporal rounding (#6456) [#6486](https://github.com/apache/datafusion-comet/pull/6486) (andygrove) +- fix: [branch-1.1] defer Parquet conversion errors until a row group is decoded, as Spark does (#6515) [#6540](https://github.com/apache/datafusion-comet/pull/6540) (andygrove) +- fix: [branch-1.1] dispatch StaticInvoke and Invoke only into Spark's own classes (#6542) [#6554](https://github.com/apache/datafusion-comet/pull/6554) (andygrove) +- fix: [branch-1.1] gate array distinct and union signed-zero semantics by Spark version (#5750) [#6561](https://github.com/apache/datafusion-comet/pull/6561) (andygrove, peterxcli) +- fix: [branch-1.1] fall back from the native Iceberg scan when a nested field was added or renamed (#6543) [#6560](https://github.com/apache/datafusion-comet/pull/6560) (andygrove) + +**Performance related:** + +- perf: optimize spark_base64 in spark-expr [#4885](https://github.com/apache/datafusion-comet/pull/4885) (andygrove) +- perf: optimize `spark_floor` (up to 4x faster) [#4911](https://github.com/apache/datafusion-comet/pull/4911) (andygrove) +- perf: cache Iceberg reflection lookups on the planning path [#5222](https://github.com/apache/datafusion-comet/pull/5222) (andygrove) +- perf: vectorize integer-to-decimal cast [#4939](https://github.com/apache/datafusion-comet/pull/4939) (andygrove) +- perf: compute spark_size list lengths with Arrow length kernel [#5233](https://github.com/apache/datafusion-comet/pull/5233) (0lai0) +- perf: make CometShuffleBenchmark completable and fast [#5388](https://github.com/apache/datafusion-comet/pull/5388) (andygrove) +- perf: vectorize Map in spark_size via offset buffer [#5395](https://github.com/apache/datafusion-comet/pull/5395) (0lai0) +- perf: improve ArrowWriter performance for fixed-length vectors [#5046](https://github.com/apache/datafusion-comet/pull/5046) (peterxcli) +- perf: reuse Arrow IPC compression context across shuffle blocks [#5038](https://github.com/apache/datafusion-comet/pull/5038) (peterxcli) +- perf: use Arrow cast for decimal rescale check [#5440](https://github.com/apache/datafusion-comet/pull/5440) (peterxcli) +- perf: bulk copy fixed-width columns in ArrowWriter [#5442](https://github.com/apache/datafusion-comet/pull/5442) (peterxcli) +- perf: serialize Python input directly from Comet Arrow vectors [#5368](https://github.com/apache/datafusion-comet/pull/5368) (sunchao) +- perf: reuse per-partition scratch in the shuffle write path [#5568](https://github.com/apache/datafusion-comet/pull/5568) (dwsmith1983) +- perf: validate shuffle IPC context reuse savings [#5727](https://github.com/apache/datafusion-comet/pull/5727) (peterxcli) +- perf: slice the child instead of gathering it when unnesting [#5667](https://github.com/apache/datafusion-comet/pull/5667) (andygrove) +- perf: evaluate posexplode array expressions once per batch [#5737](https://github.com/apache/datafusion-comet/pull/5737) (rich7420) +- perf: cache expected schemas for remote shuffle decoding [#5722](https://github.com/apache/datafusion-comet/pull/5722) (pingzh) +- perf: use Arrow cast for date to timestamp NTZ [#5735](https://github.com/apache/datafusion-comet/pull/5735) (peterxcli) +- perf: reduce allocations when collecting cache statistics [#5734](https://github.com/apache/datafusion-comet/pull/5734) (peterxcli) +- perf: avoid repeated decimal promotion in expression serialization [#5736](https://github.com/apache/datafusion-comet/pull/5736) (peterxcli) +- perf: give collect_list and collect_set a native GroupsAccumulator [#5803](https://github.com/apache/datafusion-comet/pull/5803) (andygrove) +- perf: vectorize the native map lookup behind element_at and GetMapValue [#5806](https://github.com/apache/datafusion-comet/pull/5806) (andygrove) +- perf: compile user regex patterns once per planned expression [#5612](https://github.com/apache/datafusion-comet/pull/5612) (dwsmith1983) +- perf: batch the nested-element list hash for flat struct elements [#5778](https://github.com/apache/datafusion-comet/pull/5778) (viirya) +- perf: optimize `map_sort` singleton normalization (18x faster) [#5887](https://github.com/apache/datafusion-comet/pull/5887) (viirya) +- perf: optimize `map_sort` for multi-entry string maps (up to 3x faster) [#5901](https://github.com/apache/datafusion-comet/pull/5901) (viirya) +- perf: skip calendar reconstruction in hour/minute/second and dayofweek/weekday [#5771](https://github.com/apache/datafusion-comet/pull/5771) (peterxcli) +- perf: reuse zstd compression contexts across shuffle blocks [#5565](https://github.com/apache/datafusion-comet/pull/5565) (dwsmith1983) +- perf: cache parsed plan data across a stage tasks [#5615](https://github.com/apache/datafusion-comet/pull/5615) (dwsmith1983) +- perf: spill every shuffle partition of a task into one file [#5916](https://github.com/apache/datafusion-comet/pull/5916) (peterxcli) +- perf: decode shuffle blocks against a cached schema instead of re-parsing per block [#5809](https://github.com/apache/datafusion-comet/pull/5809) (peterxcli) +- perf: project cached batches by buffer selection, prune on collated strings [#5543](https://github.com/apache/datafusion-comet/pull/5543) (andygrove) +- perf: make one thread-local access per tracked allocation [#6166](https://github.com/apache/datafusion-comet/pull/6166) (andygrove) +- perf: reuse quantile summary buffers during merge [#4932](https://github.com/apache/datafusion-comet/pull/4932) (peterxcli) +- perf: optimize list_extract without defaults using Arrow take [#5174](https://github.com/apache/datafusion-comet/pull/5174) (peterxcli) +- perf: track decimal overflow without rescanning results [#5044](https://github.com/apache/datafusion-comet/pull/5044) (peterxcli) + +**Implemented enhancements:** + +- feat: support `explode_outer` [#5192](https://github.com/apache/datafusion-comet/pull/5192) (comphead) +- feat: support timestampadd and timestampdiff via codegen dispatch [#5030](https://github.com/apache/datafusion-comet/pull/5030) (andygrove) +- feat: support \_metadata constant columns in native Parquet scan [#5237](https://github.com/apache/datafusion-comet/pull/5237) (mbutrovich) +- feat: add `make_interval` support (codegen dispatch + native) [#5039](https://github.com/apache/datafusion-comet/pull/5039) (peterxcli) +- feat: Optionally split the Iceberg V2 write operator into distinct writer and committer operations [#4658](https://github.com/apache/datafusion-comet/pull/4658) (jordepic) +- feat: detect Iceberg V2 writes and emit fall-back reasons [#5298](https://github.com/apache/datafusion-comet/pull/5298) (jordepic) +- feat: remove native cast from boolean to decimal [#5185](https://github.com/apache/datafusion-comet/pull/5185) (andygrove) +- feat: add micro benchmark runner and EC2 guide [#5374](https://github.com/apache/datafusion-comet/pull/5374) (andygrove) +- feat: build gate + inert wiring for contrib Delta scans [Delta contrib split, part 2] [#4952](https://github.com/apache/datafusion-comet/pull/4952) (schenksj) +- feat: Support native scans with unprojected Spark 4 VARIANT columns [#5377](https://github.com/apache/datafusion-comet/pull/5377) (sunchao) +- feat: expose native aggregate spill and memory metrics [#5423](https://github.com/apache/datafusion-comet/pull/5423) (sunchao) +- feat: support `WindowGroupLimitExec` [#4870](https://github.com/apache/datafusion-comet/pull/4870) (comphead) +- feat: add RSS partition writer and task-owned JNI callback (1/n) [#5473](https://github.com/apache/datafusion-comet/pull/5473) (pingzh) +- feat: native `spark_unbase64` kernel [#5451](https://github.com/apache/datafusion-comet/pull/5451) (0lai0) +- feat: add partition writer destinations to native shuffle plans (2/n) [#5476](https://github.com/apache/datafusion-comet/pull/5476) (pingzh) +- feat: add destination-aware native shuffle execution (3/n) [#5481](https://github.com/apache/datafusion-comet/pull/5481) (pingzh) +- feat: bind task-owned RSS callbacks to native shuffle plans (4/n) [#5491](https://github.com/apache/datafusion-comet/pull/5491) (pingzh) +- feat: add Celeborn shuffle manager and partition pusher (5/n) [#5501](https://github.com/apache/datafusion-comet/pull/5501) (pingzh) +- feat: complete Celeborn map-side shuffle push lifecycle (6/n) [#5513](https://github.com/apache/datafusion-comet/pull/5513) (pingzh) +- feat: add experimental native support for in-memory cache, disabled by default [#5051](https://github.com/apache/datafusion-comet/pull/5051) (andygrove) +- feat: add raw Celeborn native shuffle reader (7/n) [#5531](https://github.com/apache/datafusion-comet/pull/5531) (pingzh) +- feat: support Map for `CreateArray` literal [#5452](https://github.com/apache/datafusion-comet/pull/5452) (comphead) +- feat: route `round` on float/double through the codegen dispatcher [#5600](https://github.com/apache/datafusion-comet/pull/5600) (andygrove) +- feat: enable native-only Celeborn shuffle planning (8/n) [#5537](https://github.com/apache/datafusion-comet/pull/5537) (pingzh) +- feat: support unicode case sensitive field names for reading parquet [#5602](https://github.com/apache/datafusion-comet/pull/5602) (comphead) +- feat: implement native Iceberg V2 writer via iceberg-rust [#5361](https://github.com/apache/datafusion-comet/pull/5361) (jordepic) +- feat: Support Iceberg system functions (bucket, truncate, years/months/days/hours) natively [#5638](https://github.com/apache/datafusion-comet/pull/5638) (andygrove) +- feat: implement regr_slope, regr_intercept, regr_r2, regr_sxx, regr_syy, regr_sxy aggregates [#4775](https://github.com/apache/datafusion-comet/pull/4775) (andygrove) +- feat: carry VariantType identity through schema serialization [#5631](https://github.com/apache/datafusion-comet/pull/5631) (peterxcli) +- feat: route unrecognized StaticInvoke and Invoke through the codegen dispatcher [#5692](https://github.com/apache/datafusion-comet/pull/5692) (andygrove) +- feat: support S3 compliant filesystems [#5314](https://github.com/apache/datafusion-comet/pull/5314) (comphead) +- feat: add native spark_sequence kernel for integral element types [#5614](https://github.com/apache/datafusion-comet/pull/5614) (0lai0) +- feat: support nested types as native shuffle hash partitioning keys [#5567](https://github.com/apache/datafusion-comet/pull/5567) (viirya) +- feat: route next_day and levenshtein collated input through the codegen dispatcher [#5720](https://github.com/apache/datafusion-comet/pull/5720) (0lai0) +- chore: add_benches hash function aggregators [#5730](https://github.com/apache/datafusion-comet/pull/5730) (coderfender) +- test: cover Spark-to-Arrow batch conversion edge cases [#5713](https://github.com/apache/datafusion-comet/pull/5713) (peterxcli) +- feat: normalize marked Variant arrays at the native Parquet boundary [#5715](https://github.com/apache/datafusion-comet/pull/5715) (peterxcli) +- feat: support native concat_ws with string arrays [#5725](https://github.com/apache/datafusion-comet/pull/5725) (peterxcli) +- feat: native dynamic filter pushdown for hash join into Parquet scans [#5699](https://github.com/apache/datafusion-comet/pull/5699) (pingzh) +- feat: enable codegen dispatch for `lpad` and `rpad` [#5764](https://github.com/apache/datafusion-comet/pull/5764) (Satyr09) +- feat: address remaining issues for `CreateArray` [#5766](https://github.com/apache/datafusion-comet/pull/5766) (comphead) +- ci: stop non-gating labels from skipping Preflight [#5784](https://github.com/apache/datafusion-comet/pull/5784) (comphead) +- feat: route `abs` on interval types through the codegen dispatcher [#5622](https://github.com/apache/datafusion-comet/pull/5622) (kazantsev-maksim) +- ci: shard Iceberg Spark tests across four runners [#5459](https://github.com/apache/datafusion-comet/pull/5459) (sunchao) +- test: expand ANSI coverage for round, conv and elt [#5799](https://github.com/apache/datafusion-comet/pull/5799) (rich7420) +- test: strengthen ANSI exception assertions [#5800](https://github.com/apache/datafusion-comet/pull/5800) (rich7420) +- chore: Improve network retry configuration for maven and artifact upload [#5782](https://github.com/apache/datafusion-comet/pull/5782) (comphead) +- ci: gate the Delta contrib build on symbols, not on libcomet size [#5827](https://github.com/apache/datafusion-comet/pull/5827) (andygrove) +- test: add helpers to assert whether an expression ran natively or via codegen dispatch [#5610](https://github.com/apache/datafusion-comet/pull/5610) (andygrove) +- feat: support Spark encode expression via codegen dispatch [#5037](https://github.com/apache/datafusion-comet/pull/5037) (andygrove) +- feat: support native aggregate function `mode` [#4782](https://github.com/apache/datafusion-comet/pull/4782) (andygrove) +- test: name the whole dispatched subtree in the decimal promotion assertion [#5849](https://github.com/apache/datafusion-comet/pull/5849) (andygrove) +- feat: route translate and to_csv through codegen dispatch by default [#5032](https://github.com/apache/datafusion-comet/pull/5032) (andygrove) +- feat: expose native Parquet scan I/O and read-amplification metrics [#5453](https://github.com/apache/datafusion-comet/pull/5453) (sunchao) +- ci: move the job routing policy out of ci.yml expressions and into compute-changes.py [#5850](https://github.com/apache/datafusion-comet/pull/5850) (andygrove) +- feat: support max_by and min_by aggregate expressions [#4817](https://github.com/apache/datafusion-comet/pull/4817) (andygrove) +- ci: add a Required Checks aggregator job so main can have a required status check [#5842](https://github.com/apache/datafusion-comet/pull/5842) (andygrove) +- feat: adapt Parquet storage for Variant projection [#5794](https://github.com/apache/datafusion-comet/pull/5794) (peterxcli) +- ci: move the Spark 3.4/3.5/4.0, Iceberg, macOS and benchmark suites behind a merge queue [#5843](https://github.com/apache/datafusion-comet/pull/5843) (andygrove) +- feat(iceberg): Iceberg table format V3: apply deletion vector on reads [#5853](https://github.com/apache/datafusion-comet/pull/5853) (mbutrovich) +- test: restore ANSI array access error coverage [#5798](https://github.com/apache/datafusion-comet/pull/5798) (rich7420) +- test: cover native memory accounting boundaries [#5856](https://github.com/apache/datafusion-comet/pull/5856) (rich7420) +- test: cover date maps in Parquet temporal fuzz tests [#5877](https://github.com/apache/datafusion-comet/pull/5877) (rich7420) +- feat: Enable adaptive partial aggregation for eligible native shuffle plans [#5723](https://github.com/apache/datafusion-comet/pull/5723) (peterxcli) +- ci: move the Spark 4.1 sql_hive shards behind the merge queue [#5871](https://github.com/apache/datafusion-comet/pull/5871) (andygrove) +- ci: retry the Maven wrapper bootstrap in every job that calls ./mvnw directly [#5852](https://github.com/apache/datafusion-comet/pull/5852) (andygrove) +- test: cover slice over expression-produced non-null element arrays (#… [#5839](https://github.com/apache/datafusion-comet/pull/5839) (sam-1112) +- test: run the libhdfs suite manually instead of in CI [#5892](https://github.com/apache/datafusion-comet/pull/5892) (andygrove) +- feat: Add Lance contrib build gate [#5728](https://github.com/apache/datafusion-comet/pull/5728) (wirybeaver) +- chore: drop support for JDK 11 [#5897](https://github.com/apache/datafusion-comet/pull/5897) (manuzhang) +- test: cover regex expression routing configurations [#5917](https://github.com/apache/datafusion-comet/pull/5917) (rich7420) +- test: cover round expression routing configurations [#5878](https://github.com/apache/datafusion-comet/pull/5878) (rich7420) +- test: cover to_csv routing configurations [#5921](https://github.com/apache/datafusion-comet/pull/5921) (rich7420) +- test: cover interval expression routing configurations [#5920](https://github.com/apache/datafusion-comet/pull/5920) (rich7420) +- test: cover array and map expression routing configurations [#5918](https://github.com/apache/datafusion-comet/pull/5918) (rich7420) +- test: make the cancelled Iceberg abort test yield deterministically [#5919](https://github.com/apache/datafusion-comet/pull/5919) (andygrove) +- chore: deprecate Spark 3.4 rather than removing it in 1.1.0 [#5885](https://github.com/apache/datafusion-comet/pull/5885) (andygrove) +- test: restore the TPC-H suite’s 2 GiB off-heap budget [#5904](https://github.com/apache/datafusion-comet/pull/5904) (ErikBPF) +- test: cover string expression routing and native opt-in [#5915](https://github.com/apache/datafusion-comet/pull/5915) (rich7420) +- chore: run only the cache-writing jobs on push to main [#5930](https://github.com/apache/datafusion-comet/pull/5930) (andygrove) +- ci: publish nightly SNAPSHOT jars to repository.apache.org [#5902](https://github.com/apache/datafusion-comet/pull/5902) (andygrove) +- ci: fold the Delta build gate and PyArrow UDF suite into the merge queue tiers [#5926](https://github.com/apache/datafusion-comet/pull/5926) (andygrove) +- refactor: use Arrow decimal precision validation [#5160](https://github.com/apache/datafusion-comet/pull/5160) (Smallfu666) +- ci: shrink the pull request tier to the Linux build on the default Spark profile [#5939](https://github.com/apache/datafusion-comet/pull/5939) (andygrove) +- ci: run the non-default Spark and Iceberg suites nightly instead of in the merge queue [#5963](https://github.com/apache/datafusion-comet/pull/5963) (andygrove) +- test: cover lpad and rpad routing configurations [#5895](https://github.com/apache/datafusion-comet/pull/5895) (rich7420) +- test: preserve Spark SQL baselines for ordering-sensitive fixtures [#5914](https://github.com/apache/datafusion-comet/pull/5914) (rich7420) +- test: cover json_array_length routing configurations [#5940](https://github.com/apache/datafusion-comet/pull/5940) (peterxcli) +- test: cover SQL fixture metadata and statement parsing [#5941](https://github.com/apache/datafusion-comet/pull/5941) (peterxcli) +- test: cover crc32 on binary inputs [#5942](https://github.com/apache/datafusion-comet/pull/5942) (peterxcli) +- test: cover datetime and timezone routing configurations [#5950](https://github.com/apache/datafusion-comet/pull/5950) (rich7420) +- test: cover collation and predicate routing configurations [#5951](https://github.com/apache/datafusion-comet/pull/5951) (rich7420) +- test: cover array_intersect routing configurations [#5952](https://github.com/apache/datafusion-comet/pull/5952) (rich7420) +- test: guard native plan equality against omitted parameters [#5953](https://github.com/apache/datafusion-comet/pull/5953) (rich7420) +- test: expand replace compatibility regression coverage [#5409](https://github.com/apache/datafusion-comet/pull/5409) (sam-1112) +- feat: add native allocation accounting for memory observability [#5934](https://github.com/apache/datafusion-comet/pull/5934) (andygrove) +- feat: run rlike natively by default for Java-equivalent literal patterns [#5415](https://github.com/apache/datafusion-comet/pull/5415) (sam-1112) +- ci: write large actions/cache entries only on push to main [#5973](https://github.com/apache/datafusion-comet/pull/5973) (andygrove) +- refactor: extract shared runtime filter components [#5937](https://github.com/apache/datafusion-comet/pull/5937) (pingzh) +- test: skip the Spark RocksDB state-store suites when Comet is enabled [#5987](https://github.com/apache/datafusion-comet/pull/5987) (comphead) +- chore: bench nondeterministic and json kernels [#5989](https://github.com/apache/datafusion-comet/pull/5989) (coderfender) +- feat: add BinaryType support for SortMergeJoin [#5928](https://github.com/apache/datafusion-comet/pull/5928) (xhumanoid) +- refactor: replace remaining hand-rolled loops with Arrow kernels [#5367](https://github.com/apache/datafusion-comet/pull/5367) (0lai0) +- feat: support Spark 4 EmptyRelationExec as a native input [#5821](https://github.com/apache/datafusion-comet/pull/5821) (jeffw13) +- chore: Bench agg (welford) stats [#5988](https://github.com/apache/datafusion-comet/pull/5988) (coderfender) +- ci: rename the umbrella workflow from CI to Comet CI [#6003](https://github.com/apache/datafusion-comet/pull/6003) (andygrove) +- ci: bootstrap Maven from the setup actions so every mvnw job is covered [#5881](https://github.com/apache/datafusion-comet/pull/5881) (andygrove) +- ci: add dev/local-ci.sh to run the Spark SQL and Iceberg suites locally [#5974](https://github.com/apache/datafusion-comet/pull/5974) (comphead) +- feat: narrow strict floating-point admission for scalar sort keys [#5981](https://github.com/apache/datafusion-comet/pull/5981) (0lai0) +- feat: route compatible xxhash64 args through SparkXxhash64 [#5960](https://github.com/apache/datafusion-comet/pull/5960) (sam-1112) +- feat: support string arrays in Spark-to-Comet conversion [#5954](https://github.com/apache/datafusion-comet/pull/5954) (rich7420) +- feat: run length on binary input natively [#5874](https://github.com/apache/datafusion-comet/pull/5874) (dwsmith1983) +- chore: drop the redundant width_bucket shim registrations and guard serde uniqueness [#5873](https://github.com/apache/datafusion-comet/pull/5873) (dwsmith1983) +- feat: hook native Parquet writes into Spark's WriteFilesExec seam on Spark 4.0+ [#5763](https://github.com/apache/datafusion-comet/pull/5763) (andygrove) +- feat(iceberg): report native Iceberg scan planning metrics and scan time in the Spark UI [#6027](https://github.com/apache/datafusion-comet/pull/6027) (parthchandra) +- test: guard page skipping in the native Iceberg scan [#6040](https://github.com/apache/datafusion-comet/pull/6040) (dwsmith1983) +- test: stop claiming Spark agrees on signed-zero array literals [#6055](https://github.com/apache/datafusion-comet/pull/6055) (cestercian) +- feat: trace Arrow memory held on the JVM side [#6048](https://github.com/apache/datafusion-comet/pull/6048) (andygrove) +- feat: unix_timestamp codegen dispatch for strings, and fix pre-epoch fractional rounding [#5789](https://github.com/apache/datafusion-comet/pull/5789) (Satyr09) +- chore: remove JaCoCo from the build [#6084](https://github.com/apache/datafusion-comet/pull/6084) (andygrove) +- refactor: compose CometScanRule and CometExecRule into a single CometRule [#6082](https://github.com/apache/datafusion-comet/pull/6082) (andygrove) +- chore: add a native Iceberg write benchmark [#6038](https://github.com/apache/datafusion-comet/pull/6038) (0lai0) +- ci: retry the Spark test pre-compile step on resolution failures [#6079](https://github.com/apache/datafusion-comet/pull/6079) (andygrove) +- ci: focus Miri on unsafe row and hash tests [#6072](https://github.com/apache/datafusion-comet/pull/6072) (rich7420) +- test: run the Iceberg Spark tests with the native Iceberg writer enabled [#5677](https://github.com/apache/datafusion-comet/pull/5677) (andygrove) +- feat: support listagg / string_agg aggregate (Spark 4.0+) [#4816](https://github.com/apache/datafusion-comet/pull/4816) (andygrove) +- ci: run label-triggered CI as a separate workflow [#6161](https://github.com/apache/datafusion-comet/pull/6161) (andygrove) +- feat: always count native allocations and log executor native memory usage [#6162](https://github.com/apache/datafusion-comet/pull/6162) (andygrove) +- feat: use DataFusion unnest_outer instead of Comet ListEmptyToNullExpr [#6132](https://github.com/apache/datafusion-comet/pull/6132) (comphead) +- chore: recommend --force-with-lease when updating PR branches [#6171](https://github.com/apache/datafusion-comet/pull/6171) (mizulun) +- feat: positional round robin shuffle keyed on a row ordinal [#6095](https://github.com/apache/datafusion-comet/pull/6095) (andygrove) +- chore: deprecate spark.comet.exec.memoryPool.fraction [#6163](https://github.com/apache/datafusion-comet/pull/6163) (andygrove) +- feat: add spark.comet.explain.planOnly.enabled [#5394](https://github.com/apache/datafusion-comet/pull/5394) (andygrove) +- feat: support direct Variant projection in native Parquet scans [#5868](https://github.com/apache/datafusion-comet/pull/5868) (peterxcli) +- test: cover JSON and cast expression routing [#6068](https://github.com/apache/datafusion-comet/pull/6068) (rich7420) +- feat: remove memory accounting from Comet's on-heap mode [#6066](https://github.com/apache/datafusion-comet/pull/6066) (andygrove) +- test: accept CometHashAggregateExec in the Spark 4.0 CollationSuite hash agg check [#6220](https://github.com/apache/datafusion-comet/pull/6220) (andygrove) +- feat: per-location credentials for the S3 credential SPI [#6031](https://github.com/apache/datafusion-comet/pull/6031) (snmvaughan) +- refactor: centralize data type support predicates [#5025](https://github.com/apache/datafusion-comet/pull/5025) (peterxcli) +- chore: [branch-1.1] change version from 1.1.0-SNAPSHOT to 1.1.0 [#6236](https://github.com/apache/datafusion-comet/pull/6236) (andygrove) +- feat: [branch-1.1] built-in S3 credential provider adapters for the native Parquet scan (#6023) [#6318](https://github.com/apache/datafusion-comet/pull/6318) (andygrove, parthchandra) +- chore: [branch-1.1] Add 1.1.0 changelog [#6282](https://github.com/apache/datafusion-comet/pull/6282) (andygrove) +- test: [branch-1.1] fix flaky mixed field id directory test on Spark 3.4 and 3.5 (#6312) [#6527](https://github.com/apache/datafusion-comet/pull/6527) (andygrove) +- test: [branch-1.1] cover sliced boolean arrays in explode (#6473) [#6492](https://github.com/apache/datafusion-comet/pull/6492) (andygrove) +- test: [branch-1.1] backport the Iceberg write report and the mid-write retry test (#6155, #6111) [#6307](https://github.com/apache/datafusion-comet/pull/6307) (andygrove, sam-1112) +- feat: [branch-1.1] reuse S3 credentials until shortly before their reported expiry (#6509) [#6545](https://github.com/apache/datafusion-comet/pull/6545) (andygrove, snmvaughan) + +**Documentation updates:** + +- docs: add note about run-iceberg-tests label in CI [#5247](https://github.com/apache/datafusion-comet/pull/5247) (mbutrovich) +- docs: add 1.0.0 TPC-DS benchmark results, remove TPC-H [#5284](https://github.com/apache/datafusion-comet/pull/5284) (mbutrovich) +- docs: add 1.0.0 changelog to main [#5263](https://github.com/apache/datafusion-comet/pull/5263) (andygrove) +- docs: correct Spark 4.2 version and CI test status in installation guide [#5315](https://github.com/apache/datafusion-comet/pull/5315) (andygrove) +- docs: add suggest-native-expression skill for assessing native expression candidates [#5348](https://github.com/apache/datafusion-comet/pull/5348) (andygrove) +- docs: document C2R cost for wide/nested schemas in tuning guide [#5458](https://github.com/apache/datafusion-comet/pull/5458) (DebadityaHait) +- docs: update stale interval multiplication note in expressions.md [#5518](https://github.com/apache/datafusion-comet/pull/5518) (peterxcli) +- docs: add contributing guide link [#5521](https://github.com/apache/datafusion-comet/pull/5521) (Dharan-K) +- doc: fix benchmark examples for MacOS [#5522](https://github.com/apache/datafusion-comet/pull/5522) (xhumanoid) +- docs: add comet meeting link [#5621](https://github.com/apache/datafusion-comet/pull/5621) (coderfender) +- docs: add a contributor guide page for CI and the merge queue [#5863](https://github.com/apache/datafusion-comet/pull/5863) (andygrove) +- docs: add contributor guide page on memory management [#5933](https://github.com/apache/datafusion-comet/pull/5933) (andygrove) +- docs: how to check the scheduled CI runs are actually running [#6000](https://github.com/apache/datafusion-comet/pull/6000) (andygrove) +- docs: explain allocator hazards and diagram where memory is allocated [#6014](https://github.com/apache/datafusion-comet/pull/6014) (andygrove) +- docs: pre-render mermaid diagrams to SVG at build time [#6021](https://github.com/apache/datafusion-comet/pull/6021) (andygrove) +- docs: split the PR review skill by area and correct the shuffle contributor docs [#6018](https://github.com/apache/datafusion-comet/pull/6018) (andygrove) +- docs: correct the stale range-partitioning strict floating-point rule [#6049](https://github.com/apache/datafusion-comet/pull/6049) (andygrove) +- docs: recommend setting spark.executor.memoryOverhead alongside off-heap memory [#6051](https://github.com/apache/datafusion-comet/pull/6051) (andygrove) +- docs: exempt testing-category and internal configs from the versioning policy [#6089](https://github.com/apache/datafusion-comet/pull/6089) (andygrove) +- docs: make doc changes weekly comet sync [#6160](https://github.com/apache/datafusion-comet/pull/6160) (coderfender) +- docs: correct the plugin and shuffle sections of the plugin overview [#6197](https://github.com/apache/datafusion-comet/pull/6197) (andygrove) +- docs: fix stale and missing native Iceberg write details [#6150](https://github.com/apache/datafusion-comet/pull/6150) (andygrove) +- docs: add an Iceberg writes contributor guide and review skill [#6149](https://github.com/apache/datafusion-comet/pull/6149) (andygrove) +- docs: correct which operators run separate native plans in a task [#6194](https://github.com/apache/datafusion-comet/pull/6194) (andygrove) +- docs: update the user guide for the 1.1.0 release [#6168](https://github.com/apache/datafusion-comet/pull/6168) (andygrove) +- docs: [branch-1.1] backport the 1.1.0 upgrade notes and tuning guide updates (#6237, #6244, #6248) [#6265](https://github.com/apache/datafusion-comet/pull/6265) (andygrove) +- ci: [branch-1.1] run every tier on release-branch pull requests, and the full suite before an RC (#6218) [#6284](https://github.com/apache/datafusion-comet/pull/6284) (andygrove) +- docs: [branch-1.1] generate release docs for 1.1.0 [#6344](https://github.com/apache/datafusion-comet/pull/6344) (andygrove) +- docs: [branch-1.1] credit co-authors in the 1.1.0 changelog [#6359](https://github.com/apache/datafusion-comet/pull/6359) (andygrove) + +**Other:** + +- chore: start 1.1.0 development [#5242](https://github.com/apache/datafusion-comet/pull/5242) (andygrove) +- refactor: replace hand-coded rollup of expression fallback reasons onto operators [#5236](https://github.com/apache/datafusion-comet/pull/5236) (andygrove) +- test: rename CometCastSuite to CometNativeCastSuite [#5268](https://github.com/apache/datafusion-comet/pull/5268) (andygrove) +- refactor: use arity helper for Int to Decimal128 reinterpretation [#5193](https://github.com/apache/datafusion-comet/pull/5193) (0lai0) +- test: add guard for Iceberg version on Variant fallback test [#5278](https://github.com/apache/datafusion-comet/pull/5278) (mbutrovich) +- chore: update documentation links for 1.0.0 release [#5290](https://github.com/apache/datafusion-comet/pull/5290) (andygrove) +- chore(deps): bump reqsign-core from 3.2.0 to 3.2.1 in /native in the all-other-cargo-deps group [#5287](https://github.com/apache/datafusion-comet/pull/5287) (dependabot[bot]) +- chore(deps): bump object_store_opendal from 0.57.0 to 0.58.0 in /native [#5289](https://github.com/apache/datafusion-comet/pull/5289) (dependabot[bot]) +- chore(deps): bump the codeql-actions group with 2 updates [#5286](https://github.com/apache/datafusion-comet/pull/5286) (dependabot[bot]) +- bug: Revert "chore(deps): bump object_store_opendal from 0.57.0 to 0.58.0 in /native" [#5332](https://github.com/apache/datafusion-comet/pull/5332) (coderfender) +- test: cover the narrowing direction of cast_and_stamp_schema [#5285](https://github.com/apache/datafusion-comet/pull/5285) (andygrove) +- chore(deps): bump opendal from 0.57.0 to 0.58.1 in /native [#5324](https://github.com/apache/datafusion-comet/pull/5324) (manuzhang) +- chore: respect Cargo parallelism settings for native release builds [#5344](https://github.com/apache/datafusion-comet/pull/5344) (pingzh) +- test: make expression benchmark harness fair and reproducible [#5371](https://github.com/apache/datafusion-comet/pull/5371) (andygrove) +- chore(deps): bump the all-other-cargo-deps group in /native with 5 updates [#5360](https://github.com/apache/datafusion-comet/pull/5360) (dependabot[bot]) +- test: fix vacuous signed-zero coverage in SQL file tests [#5393](https://github.com/apache/datafusion-comet/pull/5393) (sam-1112) +- chore: fix clippy warnings for Rust 1.98 [#5400](https://github.com/apache/datafusion-comet/pull/5400) (ywskycn) +- chore(deps): bump the codeql-actions group with 2 updates [#5405](https://github.com/apache/datafusion-comet/pull/5405) (dependabot[bot]) +- test: strengthen signed-zero SQL assertions [#5404](https://github.com/apache/datafusion-comet/pull/5404) (sunchao) +- chore(deps): bump actions/checkout from 6 to 7 [#5406](https://github.com/apache/datafusion-comet/pull/5406) (dependabot[bot]) +- test: add expression fallback-invariance suite [#5329](https://github.com/apache/datafusion-comet/pull/5329) (4ktLuffy) +- refactor: delegate ANSI integer arithmetic to arrow checked kernels [#5280](https://github.com/apache/datafusion-comet/pull/5280) (kazantsev-maksim) +- test: docs and test-coverage hardening for native collect_list / collect_set [#5055](https://github.com/apache/datafusion-comet/pull/5055) (andygrove) +- test: deduplicate Throwable cause-chain traversal [#5441](https://github.com/apache/datafusion-comet/pull/5441) (peterxcli) +- refactor: use Arrow timezone type [#5129](https://github.com/apache/datafusion-comet/pull/5129) (Hashim1999164) +- chore: add md formatting to `make format` [#5460](https://github.com/apache/datafusion-comet/pull/5460) (comphead) +- test: fail unexpected vacuous fallback-invariance checks [#5417](https://github.com/apache/datafusion-comet/pull/5417) (sunchao) +- ci: cache Maven distributions and retry bootstrap downloads [#5422](https://github.com/apache/datafusion-comet/pull/5422) (sunchao) +- chore: add math benches [#5479](https://github.com/apache/datafusion-comet/pull/5479) (coderfender) +- chore: Bench string functions [#5492](https://github.com/apache/datafusion-comet/pull/5492) (coderfender) +- chore(deps): bump actions/setup-java from 5 to 6 [#5524](https://github.com/apache/datafusion-comet/pull/5524) (dependabot[bot]) +- chore(deps): bump the codeql-actions group with 2 updates [#5523](https://github.com/apache/datafusion-comet/pull/5523) (dependabot[bot]) +- chore: bench additional math function [#5520](https://github.com/apache/datafusion-comet/pull/5520) (coderfender) +- chore: remove redundant condition and stray println in shuffle code [#5562](https://github.com/apache/datafusion-comet/pull/5562) (viirya) +- test: cover struct data columns in native shuffle [#5564](https://github.com/apache/datafusion-comet/pull/5564) (viirya) +- chore: rename .claude directory to vendor-neutral .ai [#5594](https://github.com/apache/datafusion-comet/pull/5594) (andygrove) +- chore: bench additional scalar functions [#5598](https://github.com/apache/datafusion-comet/pull/5598) (coderfender) +- test: verify Celeborn reflection compatibility [#5604](https://github.com/apache/datafusion-comet/pull/5604) (pingzh) +- test: enable native path in lower/upper_enabled sql fixtures [#5619](https://github.com/apache/datafusion-comet/pull/5619) (cestercian) +- chore: add benches datetime [#5620](https://github.com/apache/datafusion-comet/pull/5620) (coderfender) +- chore: move dead and defensive serde guards out of convert [#5595](https://github.com/apache/datafusion-comet/pull/5595) (andygrove) +- test: exercise the Iceberg write split-operator plan in Iceberg's own suites [#5640](https://github.com/apache/datafusion-comet/pull/5640) (andygrove) +- chore(deps): bump the codeql-actions group with 2 updates [#5669](https://github.com/apache/datafusion-comet/pull/5669) (dependabot[bot]) +- test: add explode operator microbenchmark [#5381](https://github.com/apache/datafusion-comet/pull/5381) (andygrove) +- deps: bump DataFusion 55.0 and Arrow/Parquet 59.2 [#5262](https://github.com/apache/datafusion-comet/pull/5262) (mbutrovich) +- chore: add benches array functions [#5700](https://github.com/apache/datafusion-comet/pull/5700) (coderfender) +- test: restore Comet coverage for recursive HAVING and ORDER BY [#5755](https://github.com/apache/datafusion-comet/pull/5755) (rich7420) +- test: restore Parquet V2 writer and delta encoding coverage [#5760](https://github.com/apache/datafusion-comet/pull/5760) (rich7420) +- chore: Add benches for datetime funcs [#5767](https://github.com/apache/datafusion-comet/pull/5767) (coderfender) +- bench: add a benchmark for the Spark hash kernels [#5765](https://github.com/apache/datafusion-comet/pull/5765) (viirya) +- test: restore Spark 4.1 Variant shredding suites [#5745](https://github.com/apache/datafusion-comet/pull/5745) (rich7420) +- refactor: share one helper for pushing a struct's null mask into its children [#5769](https://github.com/apache/datafusion-comet/pull/5769) (viirya) +- ci: label pull requests by changed paths and title prefix [#5762](https://github.com/apache/datafusion-comet/pull/5762) (dwsmith1983) +- test: cover ambiguous exact nested Parquet field matches [#5751](https://github.com/apache/datafusion-comet/pull/5751) (peterxcli) +- bench: measure nested types as native shuffle hash partitioning keys [#5788](https://github.com/apache/datafusion-comet/pull/5788) (viirya) +- deps: bump to datafusion 55.1.0 [#5865](https://github.com/apache/datafusion-comet/pull/5865) (comphead) +- bench: isolate map normalization and nested key hashing [#5822](https://github.com/apache/datafusion-comet/pull/5822) (viirya) +- bench: add a shuffle read benchmark covering the per-block schema parse [#5805](https://github.com/apache/datafusion-comet/pull/5805) (peterxcli) +- chore(deps): bump actions/setup-java from 4 to 6 [#6012](https://github.com/apache/datafusion-comet/pull/6012) (dependabot[bot]) +- chore(deps): bump the codeql-actions group with 2 updates [#6011](https://github.com/apache/datafusion-comet/pull/6011) (dependabot[bot]) +- deps: bump the iceberg-rust pin to bb1e4a4 and document why it is pinned [#6094](https://github.com/apache/datafusion-comet/pull/6094) (andygrove) + +## Credits + +Thank you to everyone who contributed to this release. Here is a breakdown of commits (PRs merged) per contributor. A PR with commits from more than one person counts for each of them. + +``` + 147 Andy Grove + 62 Peter Lee + 26 Chao Sun + 26 KUAN-HAO HUANG + 19 Ping Zhang + 15 Oleks V + 14 dustin + 13 Bhargava Vadlamani + 13 Liang-Chi Hsieh + 11 dependabot[bot] + 10 ChenChen Lai + 8 sam-1112 + 6 Matt Butrovich + 5 Han-Yin Chang + 5 Parth Chandra + 4 Erik Bogado + 4 Steve Vaughan + 4 Wei Yan + 3 Jordan Epstein + 3 Manu Zhang + 2 Alexey + 2 Cestercian + 2 Daipayan Mukherjee + 2 Feng Zhang + 2 Kazantsev Maksim + 2 YuLun Mao + 2 liupoyi-1031 + 1 Dharan-K + 1 Hashim Khan + 1 Henos D + 1 Jeff Wang + 1 LinSimon-901101 + 1 Michael Taranov + 1 Scott Schenkein + 1 Stefan Wang + 1 Unik Dahal + 1 Viacheslav Inozemtsev + 1 Xuanyi Li + 1 debaditya + 1 divyank sameer shah + 1 hsiang-c +``` + +Thank you also to everyone who contributed in other ways such as filing issues, reviewing PRs, and providing feedback on this release. diff --git a/docs/source/changelog/index.md b/docs/source/changelog/index.md index e79878f7e70..792dfa2df3b 100644 --- a/docs/source/changelog/index.md +++ b/docs/source/changelog/index.md @@ -24,6 +24,7 @@ Per-release change logs for Apache DataFusion Comet. ```{toctree} :maxdepth: 1 +1.1.0 <1.1.0> 1.0.0 <1.0.0> 0.17.1 <0.17.1> 0.17.0 <0.17.0> diff --git a/docs/source/contributor-guide/adding_a_new_expression.md b/docs/source/contributor-guide/adding_a_new_expression.md index d9e7707bab4..a246ea70bb6 100644 --- a/docs/source/contributor-guide/adding_a_new_expression.md +++ b/docs/source/contributor-guide/adding_a_new_expression.md @@ -408,6 +408,7 @@ For full documentation on the test file format, including directives like `Confi - **Cover both column references and literals.** Comet often uses different code paths for each. The Comet SQL Tests suite automatically disables constant folding, so all-literal queries are evaluated natively. - **Include edge cases** such as `NULL`, empty strings, boundary values, `NaN`, and multibyte UTF-8 characters. +- **For `FLOAT` and `DOUBLE` input, include `-0.0` and a NaN with the sign bit set.** [Floating-Point Semantics](floating_point.md) explains which Spark rule an expression follows and how to test it. - **Keep one file per expression** to make failures easy to locate. - **Pin the mechanism when it depends on the input.** If your serde routes some input types natively and others through the JVM codegen dispatcher, say so with `expect_native(...)` and `expect_dispatch(...)`. A plain `query` cannot tell the two apart, so without these a later change that swaps one for the other passes silently. See [Comet SQL Tests](sql-file-tests.md). diff --git a/docs/source/contributor-guide/adding_a_new_operator.md b/docs/source/contributor-guide/adding_a_new_operator.md index 15c4e9169b6..85946f44fd9 100644 --- a/docs/source/contributor-guide/adding_a_new_operator.md +++ b/docs/source/contributor-guide/adding_a_new_operator.md @@ -148,7 +148,10 @@ own operators itself. Spark computes a typed operation's rows one at a time, as the conversion fills a whole Arrow batch first. So the rule leaves the output unconverted where a limit, a `mapPartitions` function, or code reading `Dataset.rdd` could stop reading it early, unless an operator that reads all of its input first, such as an exchange, a sort, or a hash aggregate, -sits in between. Fusing the deserializer, the `Invoke` that calls the user function, and the +sits in between. For the same reason, it leaves the output unconverted when the plan uses +`input_file_name()`, `input_file_block_start()` or `input_file_block_length()`: filling the batch +moves the scan's reader past the file that these report when Spark evaluates them above the +conversion. Fusing the deserializer, the `Invoke` that calls the user function, and the serializer of `Dataset.map` into one projection in the JVM codegen dispatcher was tried in [#5714](https://github.com/apache/datafusion-comet/pull/5714) and dropped. The dispatcher only calls into Spark's own classes, and the conversion gets nearly the same speedup for `map` while also @@ -731,6 +734,34 @@ Use `QueryPlanSerde.exprToProto` to convert Spark expressions to protobuf: val protoExpr = exprToProto(sparkExpr, inputSchema) ``` +### Restoring the Spark operator (`sparkFallback`) + +`CometExec.originalPlan` is the Spark operator this node replaced. `CometExecRule` copies +`originalPlan.logicalLink` onto the Comet node, which is how AQE finds the node again when it +re-plans a stage. `RevertNativeForTransitionHeavyStages` calls `sparkFallback(newChildren)` to +rebuild that Spark operator with the children of the reverted stage. + +The default implementation is `originalPlan.withNewChildren(newChildren)`. It refuses a null +`originalPlan`, an `originalPlan` that is one of the node's own children, or a different number of +children than the Spark operator has. + +Override `sparkFallback` when conversion changes the plan shape, so the restored node is not that +Spark operator with the same children. `CometNativeWriteExec` replaces a `DataWritingCommandExec` +and drops the `WriteFilesExec` under it; its override puts that wrapper back around the restored +input. `CometIcebergWriteExec` keeps the same shape as `IcebergWriteExec`, so the default is +enough. + +Also override `sparkFallback` when the operator's live state differs from `originalPlan`. +`CometNativeScanExec` restores its current partition and data filters, and +`CometIcebergNativeScanExec` restores its current runtime filters, so AQE's executable DPP +subqueries survive reversion. `CometLocalTopKExec` returns the restored child directly: Comet +inserted that local candidate selection, and only the outer TopK restores Spark's offset and +projection. Rebuilding the original TopK at both nodes would apply it twice. + +Do not point `originalPlan` at a child. If that child is a shuffle or query stage, the copied +logical link puts this node inside the stage's `LogicalQueryStage`. AQE then re-plans a second +copy of the operator around the one that is already there. + ### Handling Fallback Use `withInfo` to tag operators with fallback reasons: diff --git a/docs/source/contributor-guide/benchmarking_macos.md b/docs/source/contributor-guide/benchmarking_macos.md index 13ed2c26a27..4f3541811ce 100644 --- a/docs/source/contributor-guide/benchmarking_macos.md +++ b/docs/source/contributor-guide/benchmarking_macos.md @@ -154,6 +154,7 @@ $SPARK_HOME/bin/spark-submit \ --conf spark.comet.enabled=true \ --conf spark.comet.shuffle.enabled=true \ --conf spark.comet.exec.forceShuffledHashJoin=true \ + --conf spark.comet.exec.forceShuffledHashJoin.maxBuildSize=-1 \ $DF_BENCH/runners/datafusion-comet/tpcbench.py \ --benchmark tpch \ --data $BENCH_DATA/ \ diff --git a/docs/source/contributor-guide/expression-audits/array_funcs.md b/docs/source/contributor-guide/expression-audits/array_funcs.md index 2eed573d7f0..50ac45b852d 100644 --- a/docs/source/contributor-guide/expression-audits/array_funcs.md +++ b/docs/source/contributor-guide/expression-audits/array_funcs.md @@ -48,7 +48,7 @@ - Spark 3.5.8 (audited 2026-05-27): baseline. `ArrayContains(left, right) extends BinaryExpression with NullIntolerant with Predicate`; `inputTypes` uses `findWiderTypeWithoutStringPromotionForTwo`. Wired as `CometScalarFunction("array_contains")`. - Spark 4.0.1 (audited 2026-05-27): `NullIntolerant` trait replaced by `nullIntolerant: Boolean`; `checkInputDataTypes` adopts `DataTypeUtils.sameType` (collation-aware in 4.x). - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. -- Float/double arrays containing NaN and signed zero match Spark; DataFusion canonicalizes them the same way as Spark's `SQLOrderingUtil`. +- Spark 3.4, 3.5, 4.0, 4.1, 4.2 (2026-10-02, [#6520](https://github.com/apache/datafusion-comet/issues/6520)): datafusion-spark's `array_contains` compares floats by their bits, where Spark's `genEqual` treats `-0.0` as `0.0` and all NaNs as equal. Flat `FLOAT` and `DOUBLE` arrays now run through Comet's `spark_array_contains`, which applies Spark's equality, instead of the codegen dispatcher. Nested float elements stay on the dispatcher, where Spark's generated code is faster and keeps string collations, as does a value other than a literal or column read over a nullable array, which Spark skips for a null array. Tracked upstream as [apache/datafusion#25981](https://github.com/apache/datafusion/issues/25981). ## array_distinct diff --git a/docs/source/contributor-guide/expression-audits/conversion_funcs.md b/docs/source/contributor-guide/expression-audits/conversion_funcs.md index e31d5c92d9b..0b52c96d014 100644 --- a/docs/source/contributor-guide/expression-audits/conversion_funcs.md +++ b/docs/source/contributor-guide/expression-audits/conversion_funcs.md @@ -24,7 +24,7 @@ ## cast - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8 modulo `Cast.canUpCast` refactored to delegate to `UpCastRule.canUpCast`. -- Spark 3.5.8 (audited 2026-05-27): baseline. `Cast(child, dataType, timeZoneId, evalMode)`; eval modes are `LEGACY`, `ANSI`, `TRY`. The legacy `Cast.canCast` matrix and the `Cast.canAnsiCast` matrix decide acceptance per type pair. Comet routes via `CometCast` (`spark/src/main/scala/org/apache/comet/expressions/CometCast.scala`) using a per-source-type support matrix that returns `Compatible`, `Incompatible(reason)`, or `Unsupported(reason)`; literal children are short-circuited to `Compatible()` so `CometLiteral` validates them. The serialized `Cast` proto carries `datatype`, `evalMode`, `timezone` (default `UTC`), `allowIncompat` (from `spark.comet.expression.Cast.allowIncompatible`), and `isSpark4Plus`. The native side (`native/spark-expr/src/conversion_funcs/cast.rs`) implements explicit per-eval-mode branches for narrowing numeric casts that match Spark's overflow exceptions, and falls through to DataFusion `cast_with_options(safe = !ANSI)` for the rest. +- Spark 3.5.8 (audited 2026-05-27): baseline. `Cast(child, dataType, timeZoneId, evalMode)`; eval modes are `LEGACY`, `ANSI`, `TRY`. The legacy `Cast.canCast` matrix and the `Cast.canAnsiCast` matrix decide acceptance per type pair. Comet routes via `CometCast` (`spark/src/main/scala/org/apache/comet/expressions/CometCast.scala`) using a per-source-type support matrix that returns `Compatible`, `Incompatible(reason)`, or `Unsupported(reason)`; literal children are short-circuited to `Compatible()` so `CometLiteral` validates them. The serialized `Cast` proto carries `datatype`, `evalMode`, `timezone` (default `UTC`), and `isSpark4Plus`; `allowIncompatible` is enforced only by the JVM planner and is not sent to the native cast kernel. The native side (`native/spark-expr/src/conversion_funcs/cast.rs`) implements explicit per-eval-mode branches for narrowing numeric casts that match Spark's overflow exceptions, and falls through to DataFusion `cast_with_options(safe = !ANSI)` for the rest. - Spark 4.0.1 (audited 2026-05-27): `VariantType` added; `StringType` literals replaced with `_: StringType` to accommodate collated strings. `(TimestampType, ByteType|ShortType|IntegerType)` added to `canAnsiCast`. `NullIntolerant` -> `nullIntolerant: Boolean` refactor. New `ToPrettyString.BinaryFormatter` semantics for `Binary -> String` are replicated natively via `spark_binary_formatter`. Numeric-to-numeric matrix unchanged. - Spark 4.1.1 (audited 2026-05-27): `TimeType` added; many `TimeType` arms in `canCast`/`canAnsiCast`. Geospatial `GeographyType` / `GeometryType` types added with their own conversion rules. Numeric-to-numeric matrix unchanged. - Known divergences and gaps: diff --git a/docs/source/contributor-guide/expression-audits/datetime_funcs.md b/docs/source/contributor-guide/expression-audits/datetime_funcs.md index 242fd7e6c35..ba8aedd0b3e 100644 --- a/docs/source/contributor-guide/expression-audits/datetime_funcs.md +++ b/docs/source/contributor-guide/expression-audits/datetime_funcs.md @@ -33,6 +33,11 @@ - Constant-folded to a literal by Spark's `ComputeCurrentTime` rule before Comet sees the plan. +## date_trunc + +- Spark 4.2.0 and later raise `long overflow` when SECOND or MILLISECOND truncation falls below the smallest timestamp, like every other unit (SPARK-56663). Earlier versions wrap the Long subtraction. The serde tells the native kernel which Spark it runs against, so on 4.2 those two units also keep dictionary key masking for values near the lower bound. +- Performance (tuned 2026-10-06, PR [#5956](https://github.com/apache/datafusion-comet/pull/5956)): scalar timestamp truncation skips dictionary key masking for infallible fine units and safe UTC/NTZ values with substantial repetition. A values-only range check retains masking near the lower timestamp bound; non-UTC calendar paths and high-cardinality coarse batches retain masking to avoid errors and unnecessary calendar work on unused entries. About 19x faster for MICROSECOND with 65,536 keys and 32 distinct values. Matched Criterion coverage includes sparse/dense/all-NULL keys, high-cardinality YEAR batches, and used/unused overflow values. Benchmark: `benches/timestamp_trunc.rs`. + ## dayname - Spark 4.0+. Implemented natively: maps a `DateType` value to a fixed US-English abbreviated day name (`DayOfWeek.getDisplayName(TextStyle.SHORT, Locale.US)`), with no session-locale or timezone dependence. diff --git a/docs/source/contributor-guide/expression-audits/generator_funcs.md b/docs/source/contributor-guide/expression-audits/generator_funcs.md index 3023cfaa774..d34a0b61071 100644 --- a/docs/source/contributor-guide/expression-audits/generator_funcs.md +++ b/docs/source/contributor-guide/expression-audits/generator_funcs.md @@ -23,18 +23,18 @@ ## explode -- Handled at the operator level as a `GenerateExec` (`CometExplodeExec`), not via the expression serde maps, so it is not auto-detected by the function-registry checkbox logic. Compatible for array inputs; map inputs fall back ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). +- Handled at the operator level as a `GenerateExec` (`CometExplodeExec`), not via the expression serde maps, so it is not auto-detected by the function-registry checkbox logic. Compatible for array and map inputs. Maps emit key/value columns, with a position column for `posexplode`. ## explode_outer -- Same `CometExplodeExec` path as `explode`. Compatible for array inputs; empty and NULL arrays both emit one null-valued row per Spark's `outer` semantics, which the planner requests as DataFusion's `NullHandling::PreserveAndExpandEmpty`. Map inputs fall back. +- Same `CometExplodeExec` path as `explode`. Compatible for array and map inputs; empty and NULL collections both emit one null-valued row per Spark's `outer` semantics, which the planner requests as DataFusion's `NullHandling::PreserveAndExpandEmpty`. Map entries reuse the list unnest path and expand into key/value columns. ## posexplode -- Handled at the operator level as a `GenerateExec` (`CometExplodeExec`), like `explode`. Compatible for array inputs; map inputs fall back ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). +- Handled at the operator level as a `GenerateExec` (`CometExplodeExec`), like `explode`. Compatible for array and map inputs. Maps emit key/value columns, with a position column for `posexplode`. ## posexplode_outer -- Same `CometExplodeExec` path as `posexplode`. Compatible for array inputs; empty and NULL arrays both emit one row with null `pos` and null `value` per Spark's `outer` semantics, which the planner requests as DataFusion's `NullHandling::PreserveAndExpandEmpty`. +- Same `CometExplodeExec` path as `posexplode`. Compatible for array and map inputs; empty and NULL collections both emit one row with null `pos` and null generated columns per Spark's `outer` semantics, which the planner requests as DataFusion's `NullHandling::PreserveAndExpandEmpty`. [Spark Expression Support]: ../../user-guide/latest/expressions.md diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index 7027080bbf3..36328468f34 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -119,6 +119,16 @@ realigns under-aligned buffers on import ([apache/arrow-rs#10030](https://github the `realigns_under_aligned_decimal128` test in `scan.rs` guards against an arrow downgrade that would bring back the panic ([apache/arrow-rs#10028](https://github.com/apache/arrow-rs/issues/10028)). +### Errors from the JVM Producer + +Arrow Java's exported stream catches whatever the reader throws in `get_next` and hands native only its text, so +native fails the plan with a `CometNativeException` built from that text. `CometArrowStream.stream` wraps every +reader to keep the throwable itself until the task completes, and `CometExecIterator` rethrows it in place of the +native error. The task then fails with the exception Spark would have thrown, such as a `SparkArithmeticException` +from an upstream plan. For an input of Arrow-backed `ColumnarBatch`es, the first batch never takes this path: schema +reconciliation reads it on the JVM before the stream is exported, so what it throws propagates directly. A test of +this path therefore has to fail a later batch. + ### Schema Reconciliation `CometArrowStream.reconcileStreamSchema` advertises the stream's schema from the actual `CometVector` types in the diff --git a/docs/source/contributor-guide/floating_point.md b/docs/source/contributor-guide/floating_point.md new file mode 100644 index 00000000000..165b280c634 --- /dev/null +++ b/docs/source/contributor-guide/floating_point.md @@ -0,0 +1,212 @@ + + +# Floating-Point Semantics + +This page describes how Comet matches Spark's handling of `-0.0` and NaN in `FLOAT` and `DOUBLE` +values. It is aimed at contributors working on an expression or operator that compares, orders, +hashes, deduplicates or groups floating-point values, including values nested in arrays, structs +and maps. For user-facing differences from Spark, see +[Floating-point Number Comparison](../user-guide/latest/compatibility/floating-point.md). The work +to apply these rules systematically is tracked in +[#6385](https://github.com/apache/datafusion-comet/issues/6385). + +The short version: Spark has no single rule for floating-point equality. Each function inherits one +from the Java API that its implementation calls. Arrow and DataFusion follow IEEE 754 total order, +which matches none of them. A native implementation must follow the rule of the Spark function it +replaces, using the shared helpers in `native/spark-expr/src/float_semantics/`, and its tests must +include a NaN with the sign bit set. + +## Spark's rules + +| Rule | `-0.0` and `0.0` | NaN | Used by | +| ----------------------------------------------------------- | ------------------------- | ---------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| SQL ordering (`SQLOrderingUtil.compareDoubles`, `genEqual`) | Equal | All NaNs are equal, and NaN sorts above all values | Comparisons, `IN`, sorting and ranking, `min`/`max`, `greatest`/`least`, `sort_array`, `array_min`/`array_max`, `array_contains`, `array_position`, `array_remove`, and comparisons of arrays and structs | +| `NormalizeNaNAndZero`, inserted by Spark's optimizer | `-0.0` becomes `0.0` | Every NaN becomes the canonical NaN | Grouping keys, join keys and window partition keys | +| `Murmur3Hash` and `XxHash64` | Same hash | Hashed through `doubleToLongBits`, which canonicalizes NaN | `hash`, `xxhash64` and hash partitioning | +| `java.lang.Double.equals`, for boxed values | Distinct | All NaNs are equal | Java hash sets and maps of boxed values, and Spark's `OpenHashSet` and `OpenHashMap` from Spark 3.5.2 and 4.0.0 | +| Scala `==`, for boxed values (`BoxesRunTime.equals`) | Equal | A NaN equals nothing, not even another NaN | Scala sets and maps of `Any`, such as `collect_set`'s buffer before Spark 4.2 | +| `java.lang.Double.compare` | `-0.0` sorts before `0.0` | As in SQL ordering | An ascending `sort_array` of elements that cannot be null, in Spark's generated code | + +Some functions changed rules in a Spark release, so a native path has to follow the Spark version +it runs against: + +- `collect_set` keys its buffer with Scala's `==` before Spark 4.2, so `-0.0` and `0.0` are one + value and every NaN is a value of its own. Spark 4.2 normalizes NaN and `-0.0` first and keys the + buffer by the normalized bits, so all NaNs are one value too (SPARK-57298). +- Spark's `OpenHashSet`, and the `OpenHashMap` built on it, follow `Double.equals` only from Spark + 3.5.2 and 4.0.0 (SPARK-45599). In every 3.4 release and in 3.5.0 and 3.5.1 they match a key with + `==` but hash its bits (`doubleToLongBits`), so two NaNs never match, and `-0.0` and `0.0` match + only when probing from one reaches the other. `mode` and `percentile` count values in an + `OpenHashMap`. `array_distinct`, `array_union`, `array_intersect` and `array_except` look up the + elements of a flat array in an `OpenHashSet`, but treat all NaNs as one value. +- `mode` follows `Double.equals` from Spark 3.5.2 and 4.0.0 until Spark 4.2, which folds `-0.0` + into `0.0` first (SPARK-57329). +- `array_distinct` and `array_union` keep signed zeros apart in a flat array before Spark 4.2.0 + (before 3.5.2, unless probing merges them). Spark 4.2.0 normalizes their arguments in the plan + (SPARK-54918). From 4.0.5, 4.1.4 and 4.2.1 they normalize while they evaluate instead + (SPARK-59602). +- Map construction (`ArrayBasedMapBuilder`) finds duplicate keys by `Double.equals` in Spark 3.4 + and 3.5. From Spark 4.0 it normalizes each key first, unless + `spark.sql.legacy.disableMapKeyNormalization` is set + ([#6549](https://github.com/apache/datafusion-comet/issues/6549)). + +## How Arrow and DataFusion differ + +Arrow compares floats by IEEE 754 total order. `-0.0` sorts below `0.0`, NaNs compare by their bit +patterns, and a NaN with the sign bit set sorts below `-Infinity`. Arrow's sort, row format and +hash kernels all work from the same bits. DataFusion 55 folds `-0.0` into `0.0` in some kernels but +does not canonicalize NaN, and that has changed between DataFusion releases. Don't rely on it: +normalize the values, or compare them with the helpers below. + +NaNs with the sign bit set are the normal case on x86-64, not a corner case. Every NaN that +arithmetic produces at run time, such as `sqrt(-1)` or `Infinity - Infinity`, is +`0xfff8000000000000` in both Rust and the JVM on x86-64, while aarch64 produces +`0x7ff8000000000000`. Spark hides the difference through `doubleToLongBits`. A native path that +compares raw values puts those NaNs below every other value, so a query that passes on an Apple +Silicon laptop can fail on a Linux x86 cluster. + +## Where Comet applies the rules + +### The `float_semantics` module + +`native/spark-expr/src/float_semantics/` holds the per-value rules and the kernels built on them. +Its module documentation is the reference. In short: + +| Helper | Rule | Use it to | +| ---------------------------------------------------------------- | ---------------------------- | ----------------------------------------------------------------------------------------------------------- | +| `compare_floats`, `float_lt`, `float_gt` | SQL ordering | Compare values in a kernel that returns the original bits, as `array_min` and `greatest` do | +| `compare_float_arrays`, `compare_float_array_scalar` | SQL ordering | Compare a `FLOAT` or `DOUBLE` array with another or with a scalar in place, as `SparkComparison` does | +| `spark_comparator`, `spark_equality` | SQL ordering, at any depth | Compare arrays and structs in place | +| `normalize_float`, `normalize_floats`, `normalize_nested_floats` | `NormalizeNaNAndZero` | Normalize values before an Arrow kernel sorts, row-encodes, hashes or compares them | +| `NormalizeNaNAndZero`, `NormalizeNestedFloats` | `NormalizeNaNAndZero` | Wrap a key or operand expression. `wrap_if_needed` skips other types and keys that Spark already normalized | +| `canonicalize_nan` | `java.lang.Double.equals` | Key a hash set or map the way boxed Java values do | +| `compare_floats_java` | `java.lang.Double.compare` | Sort the way `java.util.Arrays.sort` sorts a primitive array | +| `hash_input` | `Murmur3Hash` and `XxHash64` | Get the value to hash in place of a float | + +Once `-0.0` is folded and every NaN is canonical, Arrow's total order agrees with Spark's SQL +ordering. That is why normalizing the inputs of an Arrow kernel works, as long as only the keys are +normalized and the output keeps the original values. + +### Keys, sorting and comparisons + +- Spark's optimizer wraps grouping, join and window partition keys in `NormalizeNaNAndZero` + (`NormalizeFloatingNumbers`). Comet serializes it like any other expression. +- `create_normalized_key_expr` in `native/core/src/execution/planner.rs` normalizes sort keys, + window order and partition keys, and the partition keys of `WindowGroupLimit`, so that Arrow's + sort orders them the Spark way. +- `spark_comparison` in `native/spark-expr/src/array_funcs/nested_comparison.rs` builds every native + `=`, `<>`, `<=>`, `<`, `<=`, `>`, `>=` and `IS DISTINCT FROM`, whichever operator evaluates it. It + compares the operands as they are, without normalized copies: `FLOAT` and `DOUBLE` with + `compare_float_arrays` and `compare_float_array_scalar`, and arrays and structs with + `spark_comparator`, or with `spark_equality` for `=`, `<>`, `<=>` and `IS DISTINCT FROM`. It + normalizes a literal while the plan is built so that it stays a literal. Operands of other types, + such as dictionary-encoded floats, are normalized and compared with Arrow's kernels. +- A scan's pushed-down data filters keep a float column compared with a literal other than NaN as a + plain `BinaryExpr` (`FloatOperands::Raw`) so that Parquet pruning still recognizes it, but only + while the reader prunes with them without filtering rows. A bloom filter probe hashes the + literal's bits, so `=` against either zero becomes `= -0.0 OR = 0.0`. With + `spark.comet.parquet.rowFilterPushdown.enabled=true` they compare in Spark's ordering like every + other comparison, because the reader drops the rows a filter rejects, and a raw column would + reject a stored NaN that Spark matches + ([#6702](https://github.com/apache/datafusion-comet/issues/6702)). Spark's own reader + keeps the two zeros apart in its dictionary and bloom filters, so it can skip a row group that + Comet reads. A test of this pruning writes out the expected rows instead of comparing them with + Spark's. +- `IN` normalizes its operands in the serde (`normalizeInOperand` in `predicates.scala`). +- `hash`, `xxhash64`, the native shuffle's hash partitioner and `approx_count_distinct` hash floats + through `hash_input`. + +### Spark version differences + +Rules that depend on the Spark version are decided in Scala. The serde either chooses the native +path or passes a flag to native code, as `Mode`'s `normalize_neg_zero` does. Where a patch release +changed the rule, check the runtime version down to the patch with +`Utils.majorMinorPatchVersion(SPARK_VERSION)`, as `ArraySetSupport` in `arrays.scala` does. + +Until an expression has a native implementation that follows its rule, report it `Incompatible` +for floating-point input so that Spark evaluates it. + +`spark.comet.exec.strictFloatingPoint=true` makes Comet fall back for floating-point operations +that can still differ from Spark. To gate such an operation, use +`SupportLevel.strictFloatingPointReason(dataType, what)`. It returns a reason only in strict mode, +and only for a type that contains a `FLOAT` or `DOUBLE`. + +## Choosing the rule for an expression + +1. Read Spark's implementation, both `eval` and `doGenCode`, in every Spark version that Comet + supports, at the latest patch release of each. What it calls decides the rule. `ordering.compare`, + `ordering.equiv`, `ctx.genComp` and `ctx.genEqual` mean SQL ordering. A `java.util.HashMap` or + another Java collection of boxed values means `Double.equals`, and so does an `OpenHashSet` or + `OpenHashMap` from Spark 3.5.2 and 4.0.0. A Scala collection of `Any`, such as + `mutable.HashSet[Any]`, means Scala's `==`. `java.util.Arrays.sort` on a primitive array means + `Double.compare`. +2. Check whether Spark's optimizer normalizes the input in the plan. If it does, Comet receives + normalized values on that version only. +3. Apply the rule at every depth. A float inside an array, struct or map follows the same rule, and + some DataFusion kernels fold `-0.0` only in a flat array. +4. Record what you found in the expression's [audit](expression-audits/index.md) and, for a + difference users can see, in the compatibility guide. + +## Guidelines + +- Don't compare floating-point values with Arrow's comparison, sort or hash kernels, or with + `total_cmp` or `partial_cmp`, where Spark uses SQL ordering, unless you normalize them first. +- Normalize keys and operands, not results. Spark returns the original bits: `greatest(-0.0D, 0.0D)` + returns `-0.0`, `max` returns whichever zero it reads first, and `array_min` returns the first of + equal elements. +- Don't write a local copy of a rule. If a helper is missing, add it to `float_semantics`. +- In a hot loop, use the `float_lt` and `float_gt` predicates. Comparing `compare_floats(a, b)` with + an ordering known only at run time was several times slower in a scan for a minimum. +- When equality on nested values can stop early, keep that exit. `spark_equality` rejects lists of + different lengths without comparing their elements, which an equality built on an ordering + comparator loses. + +## Testing + +- Use the edge values `0.0`, `-0.0`, a canonical NaN, a NaN with the sign bit set, `1.0`, `-1.0`, + `Infinity`, `-Infinity` and `NULL`, for both `FLOAT` and `DOUBLE`. +- Make the sign-bit NaN at query time. Spark's Parquet writer canonicalizes NaN, so a stored NaN + reads back canonical. Negate it in the query (`-d`), which flips the sign bit on every platform. + Arithmetic such as `sqrt(-1)` yields a sign-bit NaN only on x86-64, so a test that relies on it + passes on Apple Silicon without exercising anything. The writer keeps the sign of zero, so `-0.0` + can be stored directly. +- Read the values from a Parquet table, and also compare them with literals on both sides. The + [Comet SQL Tests](sql-file-tests.md) turn constant folding off, so `-0.0D` stays a literal with the + sign bit set, while `double('NaN')` is a cast. +- Cover every place the rule applies. For a comparison, that is `Project`, `Filter`, aggregate + arguments, `FILTER` clauses, join conditions, sort keys, generators and windows. +- Nest the values in arrays and structs, two levels deep, and make the nested NaN a sign-bit one. +- When arithmetic produces the NaN, run with ANSI mode both on and off + (`-- ConfigMatrix: spark.sql.ansi.enabled=false,true`), because division and remainder take + different native paths. +- Where the rule changed in a Spark release, gate the fixtures with `MinSparkVersion` or + `MaxSparkVersion`, and add the `run-all-spark-profiles` label to the pull request. CI pins one + patch release per Spark line, so when a patch release changed the rule, also test the routing in + Scala. +- Add the expression to `CometFloatSemanticsSuite`, which crosses the edge values with operator + contexts and with every expression and aggregate that accepts floating-point input. New + expressions go in its `expressions` or `aggregates` list. If one does not match Spark yet, add a + `KnownGap` that links the issue. Known gaps are strict: a case that starts matching Spark fails, + so the fix also removes the entry. + +Fixtures to start from include `expressions/conditional/float_comparisons.sql`, +`expressions/math/nan_divisor.sql`, `expressions/math/greatest_least_floating_point.sql`, +`expressions/aggregate/min_max_floating_point.sql`, `expressions/array/sort_array_floating_point.sql` +and `windows/nested_float_order_keys.sql`, all under `spark/src/test/resources/sql-tests/`. diff --git a/docs/source/contributor-guide/iceberg-writes.md b/docs/source/contributor-guide/iceberg-writes.md index 073b7c836d9..1ef7539329a 100644 --- a/docs/source/contributor-guide/iceberg-writes.md +++ b/docs/source/contributor-guide/iceberg-writes.md @@ -70,13 +70,13 @@ IcebergCommit driver: collect task commit messages, BatchWrite.c +- scans, projects, exchanges, sorts; inside AQE with or without the split ``` -| Component | Location | Role | -| ------------------------------------------------------------------------------- | -------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `IcebergWriteStrategy` | `spark/src/main/scala/org/apache/comet/iceberg/` | Planner strategy. Matches `AppendData`, `OverwriteByExpression`, `OverwritePartitionsDynamic`, `ReplaceData` (and Iceberg's own `ReplaceIcebergData`, which Iceberg 1.5.2 plans on Spark 3.4) whose `Write` is an Iceberg `SparkWrite`. | -| `IcebergWriteLogical` | same | Logical anchor for the writer, so AQE re-plans re-emit only the writer and not a second committer. | -| `IcebergWriteExec` | `spark/src/main/scala/org/apache/spark/sql/comet/` | JVM writer. Runs iceberg-java's `DataWriter` per task and returns the serialized `WriterCommitMessage` as one binary row. | -| `IcebergCommitExec` | same | Driver committer. A `V2CommandExec`, so `run()` is memoized and the commit happens once. | -| `IcebergReplaceDataShim`, `IcebergRefreshCacheShim`, `IcebergDriverMetricsShim` | `spark/src/main/spark-*/org/apache/comet/iceberg/` | Version differences: Spark 4.x operation-coded `ReplaceData` rows, cache refresh by name on 4.1+, driver metric reporting. | +| Component | Location | Role | +| ------------------------------------------------------------------------------------------------------------------------------------ | -------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| `IcebergWriteStrategy` | `spark/src/main/scala/org/apache/comet/iceberg/` | Planner strategy. Matches `AppendData`, `OverwriteByExpression`, `OverwritePartitionsDynamic`, `ReplaceData` (and Iceberg's own `ReplaceIcebergData`, which Iceberg 1.5.2 plans on Spark 3.4), plus Spark 3.5+ Iceberg `WriteDelta`. | +| `IcebergWriteLogical` | same | Logical anchor for the writer, so AQE re-plans re-emit only the writer and not a second committer. | +| `IcebergWriteExec` | `spark/src/main/scala/org/apache/spark/sql/comet/` | JVM writer. Runs iceberg-java's `DataWriter` or, for `WriteDelta`, `DeltaWriter` per task and returns the serialized `WriterCommitMessage` as one binary row. | +| `IcebergCommitExec` | same | Driver committer. A `V2CommandExec`, so `run()` is memoized and the commit happens once. | +| `IcebergReplaceDataShim`, `IcebergDeltaLogicalShim`, `IcebergDeltaWriterShim`, `IcebergRefreshCacheShim`, `IcebergDriverMetricsShim` | `spark/src/main/spark-*/org/apache/comet/iceberg/` | Version differences: operation-coded `ReplaceData` / `WriteDelta` rows, cache refresh by name on 4.1+, driver metric reporting. | Things to know before changing this layer: @@ -90,16 +90,19 @@ Things to know before changing this layer: - **Messages are collected per task as tasks finish** (`sparkContext.runJob` with a result handler), not with `executeCollect`. That is how a failed job still knows which tasks completed, so it can delete their files. -- **Commit-coordinator writes are not intercepted.** Iceberg's `SparkWrite` never asks for one; the - check in `buildTwoOp` is defensive. -- **What is not intercepted:** merge-on-read (`WriteDelta`), streaming writes, and CTAS/RTAS on - Spark 3.4. Those keep Spark's plan. +- **Commit-coordinator writes are not intercepted.** Iceberg's current `SparkWrite` and + position-delta write do not ask for one; the checks in `buildTwoOp` and `buildDeltaTwoOp` + are defensive. +- **WriteDelta stays JVM-backed.** Spark 3.5+ merge-on-read commands are intercepted by the split + plan, but `IcebergWriteExec` delegates their rows to Iceberg's JVM `DeltaWriter`; the native + Iceberg writer is explicitly declined. Spark 3.4 `WriteDelta` keeps Spark's plan. +- **What is not intercepted:** streaming writes and CTAS/RTAS on Spark 3.4. Those keep Spark's plan. On Spark 4.1+, `IcebergWriteSummaryShim` finds either Spark's `MergeRowsExec` or `CometMergeRowsExec` in the executed query and forwards its eight action counters to `BatchWrite.commit(messages, summary)`. Native instructions carry the context of each `Keep`; `Discard` counts a deletion and `Split` counts one update. Spark 4.2 uses last-attempt -accumulators, read through `MergeRowsMetricsShim`, so summary values come from the tasks that +accumulators, read through `IcebergSemanticMetricsShim`, so summary values come from the tasks that produced the committed output. Stock V2 writers still require the concrete Spark node, so `IcebergWriteStrategy` tags a copy of its logical MERGE query before planning and AQE. The serializer requires that tag on Spark 4.1+, leaving stock writers on Spark throughout replanning. @@ -107,9 +110,9 @@ serializer requires that tag on Spark 4.1+, leaving stock writers on Spark throu ## From `IcebergWrite` to `CometIcebergWrite` `CometExecRule` converts an `IcebergWriteExec` with the `CometIcebergNativeWrite` operator serde -when `spark.comet.write.iceberg.enabled` is on. Two arms in `CometExecRule` handle it: one unwraps -the double conversion AQE can produce when it re-fires write planning over a sub-tree that already -contains a `CometIcebergWriteExec`, and the other calls `convertToComet`. +when `spark.comet.write.iceberg.enabled` is on. A single arm in `CometExecRule` handles the +conversion by calling `convertToComet`. The converted node keeps the `IcebergWriteExec` as its +`originalPlan`, so AQE re-plans the write from that node. `CometIcebergNativeWrite.requiresNativeChildren` is `true`. The native writer consumes Arrow batches from its child over FFI, so the conversion is declined unless the child is already a Comet @@ -272,10 +275,22 @@ Points where Comet adapts iceberg-rust to match iceberg-java: ([#5776](https://github.com/apache/datafusion-comet/issues/5776)). Manifest order becomes the row order of an unordered read, so any map iteration that reaches the output needs the same care. +iceberg-rust keeps each open file writer private inside its rolling and partitioning writers, so +Comet wraps `ParquetWriterBuilder` in `MeteredParquetWriterBuilder`, whose files report what they +hold in memory. It hands `ParquetWriter` each file's `OutputFile` behind a `CountedOutput`, whose +writer counts the bytes that leave memory on their way to storage, and a file reports what it has +written less those. When they leave depends on the storage (`StorageWrites`). After every batch +`run_write_task` resizes the task's reservation to what the open files report plus the rows each +`PartitionFeed` holds, for the dictionary choice or for pacing, and a resize the pool refuses fails +the task. What that figure covers, and what it misses, is described under +[Native writers](memory_management.md#native-writers). + `FileIO` comes from `load_file_io` in `iceberg_common.rs`, shared with the native scan. It picks the storage backend from the data location's scheme and wires in Comet's S3 credential bridge when one is configured. For writes the bridge fails closed: if a configured provider cannot initialize, -the task fails rather than writing with the default credential chain. +the task fails rather than writing with the default credential chain. A storage scheme newly +supported for writes also needs its entry in `StorageWrites::for_location`, which treats an unknown +scheme as one that uploads in parts. The JVM eligibility gate fails closed before building that `FileIO` for an `s3` or `s3a` data location. It checks effective `fs.s3a.*` Hadoop settings against the six keys translated by @@ -396,16 +411,16 @@ the writer: run the write suites and the Iceberg Spark tests. The pin policy is ## Testing -| Suite | What it covers | -| ---------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `CometIcebergWriteActionSuite` | End-to-end writes through the split plan and the native writer: parity with iceberg-java, row-level DML, partition evolution, file order, cleanup on task and job failure, AQE re-planning. | -| `CometIcebergWriteDetectionSuite` | One case per eligibility rule, accepted and declined. | -| `CometIcebergSystemFunctionSuite` | Native `bucket`, `truncate`, `years`/`months`/`days`/`hours`, which keep a partitioned write's hash distribution and sort native end to end. | -| `CometIcebergRewriteActionSuite` | Iceberg's `rewrite_data_files` with the split plan and the native writer. | -| `IcebergWriteProtoTranslationSuite` | Translation of properties into `IcebergParquetWriteSettings` and the writer mode. | -| Rust tests in `iceberg_write.rs` and in `iceberg_partition_*.rs` | File rolling on the 1000-row grid, fanout order, clustered input checks, cleanup guard, manifest round trip, partition path rendering, partition values past `chrono`'s calendar. | -| Rust tests in `iceberg_dictionary.rs` | The per-column dictionary choice against answers recorded from parquet-mr, including columns either side of its cut-off. | -| `CometIcebergWriteBenchmark` | Native versus iceberg-java for unpartitioned, clustered, fanout and copy-on-write delete writes. It checks each arm's plan before timing it. | +| Suite | What it covers | +| ---------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `CometIcebergWriteActionSuite` | End-to-end writes through the split plan and the native writer: parity with iceberg-java, row-level DML, partition evolution, file order, cleanup on task and job failure, a fanout write outgrowing the memory pool, AQE re-planning. | +| `CometIcebergWriteDetectionSuite` | One case per eligibility rule, accepted and declined. | +| `CometIcebergSystemFunctionSuite` | Native `bucket`, `truncate`, `years`/`months`/`days`/`hours`, which keep a partitioned write's hash distribution and sort native end to end. | +| `CometIcebergRewriteActionSuite` | Iceberg's `rewrite_data_files` with the split plan and the native writer. | +| `IcebergWriteProtoTranslationSuite` | Translation of properties into `IcebergParquetWriteSettings` and the writer mode. | +| Rust tests in `iceberg_write.rs` and in `iceberg_partition_*.rs` | File rolling on the 1000-row grid, fanout order, clustered input checks, cleanup guard, manifest round trip, memory reservation, partition path rendering, partition values past `chrono`'s calendar. | +| Rust tests in `iceberg_dictionary.rs` | The per-column dictionary choice against answers recorded from parquet-mr, including columns either side of its cut-off. | +| `CometIcebergWriteBenchmark` | Native versus iceberg-java for unpartitioned, clustered, fanout and copy-on-write delete writes. It checks each arm's plan before timing it. | The Comet suites run against the Iceberg version each Spark profile pins in `spark/pom.xml`: 1.5.2 for Spark 3.4, 1.8.1 for 3.5, 1.10.0 for 4.0 and 4.2, and 1.11.0 for 4.1. Only the default profile @@ -467,8 +482,11 @@ Each of these has caused a bug on this path: ([#5691](https://github.com/apache/datafusion-comet/issues/5691), [#5693](https://github.com/apache/datafusion-comet/issues/5693), [#6141](https://github.com/apache/datafusion-comet/issues/6141)). -- **Plan rewrites must keep the write node.** Rules that restore Spark operators from a Comet - node's `originalPlan` have to handle the write execs, whose `originalPlan` today is their child +- **Plan rewrites must keep the write node.** `CometIcebergWriteExec.originalPlan` is the + `IcebergWriteExec` it replaced. Restoring Spark execution goes through + [`CometExec.sparkFallback`](adding_a_new_operator.md#restoring-the-spark-operator-sparkfallback), + which rebuilds that node around the reverted children. Pointing `originalPlan` at the child + drops the write, and AQE then treats the write as part of that child's stage ([#5719](https://github.com/apache/datafusion-comet/issues/5719)). - **The kill switch must still work.** Code that runs for Iceberg writes has to respect `spark.comet.enabled`, so that disabling Comet restores Spark's own plan diff --git a/docs/source/contributor-guide/index.md b/docs/source/contributor-guide/index.md index 475bb77358a..1abd8891d20 100644 --- a/docs/source/contributor-guide/index.md +++ b/docs/source/contributor-guide/index.md @@ -47,6 +47,7 @@ Development Guide Comet Plugin Overview Arrow FFI Timezone Handling +Floating-Point Semantics JVM Shuffle Native Shuffle Memory Management diff --git a/docs/source/contributor-guide/jvm_shuffle.md b/docs/source/contributor-guide/jvm_shuffle.md index 11014f34909..474764b90e2 100644 --- a/docs/source/contributor-guide/jvm_shuffle.md +++ b/docs/source/contributor-guide/jvm_shuffle.md @@ -43,8 +43,10 @@ JVM shuffle (`CometColumnarExchange`) is used instead of native shuffle (`CometE 1. **Shuffle mode is explicitly set to "jvm"**: When `spark.comet.shuffle.mode` is set to `jvm`. 2. **Child plan is not a Comet native operator**: When the child plan is a Spark row-based operator - (not a `CometPlan`), JVM shuffle is the only option since native shuffle requires columnar input - from Comet operators. + (not a `CometPlan`), JVM shuffle is used, since native shuffle requires columnar input from + Comet operators. The exception is `spark.comet.convert.shuffleInput.enabled`, which converts + the child's rows to Arrow with `CometSparkToColumnarExec` so that native shuffle can take the + shuffle instead. See [When Native Shuffle is Used](native_shuffle.md#when-native-shuffle-is-used). 3. **Unsupported partition key types**: `RangePartitioning` keys must be primitive, so a complex range key always falls back here. `HashPartitioning` keys must be primitive only by default: @@ -195,7 +197,9 @@ writes the same Arrow IPC block format as native shuffle, so direct read applies depends on Spark's memory mode. - Off-heap mode gets `CometUnifiedShuffleMemoryAllocator`, an ordinary Spark `MemoryConsumer` drawing from `spark.memory.offHeap.size`. When it cannot acquire a page the writer spills to - disk. + disk. An allocation that loses the task's entry in Spark's execution pool while it waits + ([SPARK-59444](https://issues.apache.org/jira/browse/SPARK-59444)) is retried a few times before + it counts as a page that could not be acquired. - On-heap mode gets `CometUnboundedShuffleMemoryAllocator`, which keeps no budget and so never refuses a page. Nothing bounds these allocations, and memory pressure never triggers a spill. That mode exists only so the Spark SQL tests can run against Comet. See diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 2a48baa3f18..af84e4a2786 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -136,7 +136,16 @@ allocator before reading anything into the split. **The JVM shuffle allocator is an ordinary Spark consumer.** `CometShuffleMemoryAllocator.getInstance` returns `CometUnifiedShuffleMemoryAllocator`, a Spark `MemoryConsumer` drawing from `spark.memory.offHeap.size`, so shuffle pages are arbitrated against Spark's other consumers in the -same task like any other allocation. +same task like any other allocation. A page or pointer array allocation that waits in Spark for +memory can wake up to find the task's entry gone from Spark's execution pool, because another +consumer of the task released its last bytes in the meantime, and Spark then throws a +`NoSuchElementException` ("key not found") instead of a grant +([SPARK-59444](https://issues.apache.org/jira/browse/SPARK-59444)). The allocator retries that +case, since the failed call was granted nothing, and after three attempts throws +`SparkOutOfMemoryError`, which the shuffle writers handle like any other refused page. Only the +shuffle allocator's callers are guarded this way. Spark's operators in the same task, such as its +sorters and aggregates, can still hit the exception until Spark re-registers a waiting task in +`ExecutionMemoryPool` ([apache/spark#58747](https://github.com/apache/spark/pull/58747)). Which allocator each call site uses, and who ends up charged for the bytes: @@ -205,6 +214,15 @@ the call never returns. Nothing releases them until Spark's final task cleanup, nobody can use for the rest of the task. Any caller that swallows the exception has to reconcile that grant, and the only figure available for doing so is the task-wide one above. +**A parked acquire can wake up to a missing task entry.** Spark removes the task's `memoryForTask` +entry when its balance reaches zero, and an acquire that was waiting in `lock.wait()` and wakes +afterwards throws a `NoSuchElementException` ("key not found" and the task id) instead of a grant. +`CometTaskMemoryManager` retries that one case. A missing entry means the task held nothing from +Spark at that moment, so the failed call has no partial grant to reconcile, and the retry registers +the task again and waits for its share as the first call would have. After a few attempts it returns +a zero grant, which the native side treats as a refusal and spills. JVM consumers that call +`allocatePage` directly are not covered; that gap is tracked in #6304. + **A consumer whose `spill` returns zero takes budget it can never give back.** `NativeMemoryConsumer.spill` returns `0`, so Spark can select it as a spill victim and reclaim nothing from it; it is only ever a spill trigger. The bytes it holds are real, so other consumers in @@ -375,6 +393,40 @@ Native operators reserve through DataFusion's `MemoryConsumer` / `MemoryReservat An operator that never calls `try_grow` is invisible to the pool no matter how much memory it uses. +### Native writers + +Both native writers reserve what they hold between batches through a single consumer per task, +`ParquetWriterExec[N]` or `IcebergWriteExec[N]`, resized after every batch. Neither can spill, so +when the pool refuses a resize the task fails with a `CometNativeException` whose message starts +`Additional allocation failed for` and names the consumer. That is a task failure Spark can retry. +Unreserved, the same memory would count only toward the container limit, where exceeding it kills +the executor. + +- `ParquetWriterExec` writes one file per task and reserves parquet-rs's estimate of the file's + in-progress row group (`ArrowWriter::memory_size`). The estimate counts encoded pages, encoder + buffers, dictionaries and Bloom filters. The writer keeps parquet-rs's default row-group limit of + 1Mi rows and sets no byte limit, so a wide schema can hold a large row group + ([#5304](https://github.com/apache/datafusion-comet/issues/5304)). +- The Iceberg writer keeps one file open per partition in a fanout write, so its reservation grows + with the number of partitions a task writes. iceberg-rust's `ParquetWriter` does not expose + parquet-rs's `memory_size`, only every byte the file has written, its in-progress row group's + encoded size included. So each file's writer counts the bytes that leave memory on their way to + storage, and the file reports what it has written less those. A local file writes a flushed row + group out at once, so it reports its in-progress row group and the few KiB parquet-rs buffers in + front of storage. S3 and GCS take a file in parts of at least 5 MiB, and OpenDAL holds each part + in memory until the next one is complete or the file closes, so there a file also reports what + it has flushed since its last part was uploaded. At the default row-group size + (`write.parquet.row-group-size-bytes`, 128 MiB) that is the last row group it flushed. The + reservation also covers the rows each partition holds back, first for its dictionary choice and + then until they fill the 1000-row unit the rolling writer is fed in. It does not cover what + parquet-rs holds beyond the encoded size: dictionary hash tables, unencoded dictionary indices + and buffer capacity. Native writes decline Bloom filters today, and neither figure would include + them. + +The writers register one consumer per task rather than one per open file because every consumer +registered with `fair_unified` lowers the share of every other consumer in the task. A consumer per +partition would shrink the task's other operators' shares as a fanout write widened. + ## Crossing the FFI boundary Batches move between the JVM and native over the Arrow C Data and C Stream interfaces, which are @@ -425,7 +477,8 @@ The pool tracks _declared reservations_. Container RSS counts _pages the process diverge for several structural reasons: - **Undeclared allocations.** Arrow array builders, expression kernels producing intermediate - arrays, decompression buffers, Parquet metadata structures, `object_store` request buffers, and + arrays, decompression buffers, Parquet metadata structures, the part of a native writer's buffers + its estimate misses (see [Native writers](#native-writers)), `object_store` request buffers, and tokio's own machinery all allocate without reserving. Only operators that were explicitly written to reserve show up in the pool. - **Rounding and padding.** Arrow buffers are padded to 64-byte boundaries and builders grow by @@ -596,4 +649,4 @@ A checklist for triaging an executor OOM kill: `batch_size * columns`, and wide or deeply nested schemas amplify it. 4. Check whether the operators involved can spill at all. `ShuffledHashJoin` cannot, so `spark.comet.exec.forceShuffledHashJoin=true` converts a spillable sort-merge join into one that - is not. + is not, for build sides under `spark.comet.exec.forceShuffledHashJoin.maxBuildSize`. diff --git a/docs/source/contributor-guide/native_shuffle.md b/docs/source/contributor-guide/native_shuffle.md index 61088447882..704fa2a0ea7 100644 --- a/docs/source/contributor-guide/native_shuffle.md +++ b/docs/source/contributor-guide/native_shuffle.md @@ -39,6 +39,12 @@ Compare this to JVM shuffle's data path: Comet Native (columnar) → ColumnarToRowExec → rows → JVM Shuffle → Arrow IPC → columnar ``` +When `RevertNativeForTransitionHeavyStages` restores the map stage to Spark execution, the +native exchange stays in place. Its input still needs Arrow-backed Comet vectors, even if the +restored Spark operator supports columnar output. The rule adds `CometSparkToColumnarExec` to +convert either Spark rows or Spark columnar batches to Arrow before the native shuffle consumes +them. Spark's `RowToColumnarExec` alone does not satisfy this input contract. + ## When Native Shuffle is Used Native shuffle (`CometExchange`) is selected when all of the following conditions are met: @@ -46,7 +52,13 @@ Native shuffle (`CometExchange`) is selected when all of the following condition 1. **Shuffle mode allows native**: `spark.comet.shuffle.mode` is `native` or `auto`. 2. **Child plan is a Comet native operator**: The child must be a `CometPlan` that produces - columnar output. Row-based Spark operators require JVM shuffle. + columnar output. Row-based Spark operators require JVM shuffle, unless + `spark.comet.convert.shuffleInput.enabled` is set. Then a shuffle that JVM shuffle would take + uses native shuffle instead, provided native shuffle supports the partitioning and the + columns, and a `CometSparkToColumnarExec` converts the child's rows to Arrow for it. A shuffle + that hashes a string, a value computed from a string, or a decimal wider than 18 digits stays + on JVM shuffle, because native shuffle would not put every row in the partition Spark's + partitioner does (`CometShuffleExchangeExec.convertsInputForNativeShuffle`). 3. **Supported partitioning type**: Native shuffle supports: - `HashPartitioning` @@ -160,11 +172,14 @@ The native shuffle implementation is its own workspace crate, `datafusion-comet- - The child plan's `nativeOp` directly, when `CometShuffleExchangeExec`'s child is a `CometNativeExec` subtree. The upstream operators run inside the same `CometExecIterator` as the writer, with no JVM-to-native batch boundary between them. - - A synthetic `Scan("ShuffleWriterInput")` placeholder, when the dep was built via the - convenience `prepareShuffleDependency(rdd, ...)` overload (used by - `CometCollectLimitExec` and `CometTakeOrderedAndProjectExec`, or when the - exchange's child is a non-native `CometPlan` such as `CometSparkToColumnarExec`). Native - code reads `ColumnarBatch`es from the JVM input iterator via Arrow C Stream Interface. + - A synthetic `Scan("ShuffleWriterInput")` placeholder, when the exchange's child is a + non-native `CometPlan`, and for the dependencies that `CometCollectLimitExec` and + `CometTakeOrderedAndProjectExec` build. Native code reads the input through the Arrow C + Stream Interface. A `CometNativeArrowSource` child, such as `CometSparkToColumnarExec`, + exports its own Arrow stream from `doExecuteAsArrowStream()`, which + `prepareArrowStreamShuffleDependency` hands to native. Any other input is an + `RDD[ColumnarBatch]`, which the convenience `prepareShuffleDependency(rdd, ...)` overload + wraps in a stream. 2. **Native execution**: A single `CometExecIterator` per partition runs the unified plan. diff --git a/docs/source/contributor-guide/spark-sql-tests.md b/docs/source/contributor-guide/spark-sql-tests.md index c4c39764a37..3918f69ffa5 100644 --- a/docs/source/contributor-guide/spark-sql-tests.md +++ b/docs/source/contributor-guide/spark-sql-tests.md @@ -32,7 +32,9 @@ Here is an overview of the changes that we need to make to Spark: - Modify `project/SparkBuild.scala` so that, when Comet is enabled, the `sql` and `hive` test JVMs use `CometShuffleManager` as the default shuffle manager. Comet disables itself when its shuffle manager is not registered, so without this default the suites that build their own `SparkSession` or `SparkContext` would - run without Comet + run without Comet. From Spark 3.5, where Comet's cache format is the default, the same JVMs also default + `spark.sql.cache.serializer` to Comet's `ArrowCachedBatchSerializer`, because Spark keeps the first cache + serializer it loads for the rest of the JVM Here are the steps involved in running the Spark SQL tests with Comet, using Spark 3.4.3 for this example. diff --git a/docs/source/contributor-guide/sql-file-tests.md b/docs/source/contributor-guide/sql-file-tests.md index 57cc717547f..62bff931bb1 100644 --- a/docs/source/contributor-guide/sql-file-tests.md +++ b/docs/source/contributor-guide/sql-file-tests.md @@ -98,6 +98,9 @@ Sets a Spark SQL config for all queries in the file. -- Config: spark.sql.ansi.enabled=true ``` +The first `=` separates the key from the value. Values containing `=`, such as base64 +encryption keys, are preserved. An empty value is allowed; the key must be non-empty. + #### `ConfigMatrix` Runs the entire file once per combination of values. Multiple `ConfigMatrix` lines produce a @@ -117,6 +120,47 @@ sql-file: expressions/conditional/in_set.sql [spark.sql.optimizer.inSetConversio Only add a `ConfigMatrix` directive when there is a real reason to run the test under multiple configurations. Do not add `ConfigMatrix` directives speculatively. +Values are comma-separated and must be non-empty. As with `Config`, any `=` within a +value is preserved. A matrix value overrides a `Config` value for the same key. + +#### `ExcludeRules` + +Adds fully qualified optimizer rule names to the exclusion list. Constant folding remains +disabled by default. + +```sql +-- ExcludeRules: org.apache.spark.sql.catalyst.optimizer.NullPropagation +-- ExcludeRules: org.apache.spark.sql.catalyst.optimizer.ConvertToLocalRelation +``` + +Repeated lines are combined, and duplicate rules are removed. Rules declared through +`Config: spark.sql.optimizer.excludedRules=...` are also retained; a matrix selects the +value of that config before `ExcludeRules` and the default exclusion are added. Unknown +rules and rules Spark does not permit excluding fail the fixture rather than silently +leaving the optimizer unchanged. + +Exclude `NullPropagation` when the test needs a NULL argument to reach an expression. +Disabling ConstantFolding alone does not prevent Spark from replacing `abs(NULL)` with a +NULL literal. + +#### `ConstantFolding` + +Lets Spark fold constants for a fixture that tests folded complex literals. + +```sql +-- ConstantFolding: enabled +``` + +`disabled` is the default. With `enabled`, an all-literal `array(...)`, `map(...)` or +`named_struct(...)` can become a complex `Literal` before Comet sees it. This exercises +Comet's literal handling, rather than its expression constructors. Ordinary all-literal +function queries may instead fold away entirely; use column inputs when those functions +must execute in Comet. Combining `enabled` with an explicit ConstantFolding exclusion +is rejected as contradictory. Repeated `ConstantFolding` directives must agree. + +Malformed configuration directives fail that fixture with its file path and source line +instead of being treated as ordinary comments. Other fixtures still run. + #### `MinSparkVersion` Skips the file when running on a Spark version older than the specified version. @@ -148,6 +192,34 @@ statement INSERT INTO my_table VALUES (1, 2.0), (3, 4.0), (NULL, NULL) ``` +Use `SET` statements to change SQL configs between queries: + +```sql +-- Starts with the fixture default, ANSI disabled. +query +SELECT CAST(v AS INT) FROM my_strings + +statement +SET spark.sql.ansi.enabled=true + +query expect_error(CAST_INVALID_INPUT) +SELECT CAST(v AS INT) FROM my_strings +``` + +All SQLConf changes are restored when the fixture finishes, including on failure and +between matrix combinations. Keys changed by `SET` need not also appear in `Config`. +`RESET` is restored too, but it removes the selected setting for the remaining queries; +`RESET` without a key also resets the harness's ANSI and optimizer settings. A mid-file +`SET spark.sql.optimizer.excludedRules=...` replaces the whole exclusion list, including +the default ConstantFolding exclusion. + +Enabling the codegen dispatcher through `SET` still requires a non-error sentinel when +the file contains `expect_error`, just as enabling it through `Config` does. Run that +sentinel under the same settings as the error query. + +This isolation covers SQL configs, not `USE`, temporary views or functions, or SQL +session variables. Manage those separately. + ### Queries A `query` block executes a SELECT and compares results between Spark and Comet. The query @@ -345,7 +417,9 @@ planning, so Comet would never see them. However, `CometSqlFileTestSuite` automa disables constant folding (by excluding `ConstantFolding` from the optimizer rules), so all-literal queries are evaluated by Comet's native engine. This means you can use the default `query` mode for all-literal cases and they will be tested natively just like -column-based queries. +column-based queries, subject to other optimizer rewrites such as `NullPropagation` and +Comet's own plan-time literal handling. `ConstantFolding: enabled` opts out of the default +exclusion for fixtures that deliberately test folded literals. #### Cover edge cases diff --git a/docs/source/user-guide/index.md b/docs/source/user-guide/index.md index 85b129a4359..c6328493dd9 100644 --- a/docs/source/user-guide/index.md +++ b/docs/source/user-guide/index.md @@ -22,7 +22,7 @@ under the License. The Comet user guide covers installation, configuration, supported data sources, supported operators and expressions, and tuning advice for running Apache Spark with Comet acceleration. -The **current stable release** is 1.0.x. Most users should start there. The **development snapshot** +The **current stable release** is 1.1.x. Most users should start there. The **development snapshot** tracks the upcoming release and may include features and fixes that are not yet generally available. User guides for **older releases** are kept for reference under [Older Versions](older-versions). @@ -31,7 +31,7 @@ User guides for **older releases** are kept for reference under [Older Versions] :caption: User Guide :hidden: -1.0.x (current) <1.0/index> +1.1.x (current) <1.1/index> Development Snapshot Older Versions ``` diff --git a/docs/source/user-guide/latest/compatibility/floating-point.md b/docs/source/user-guide/latest/compatibility/floating-point.md index 455f58c3b63..1324dbf87fd 100644 --- a/docs/source/user-guide/latest/compatibility/floating-point.md +++ b/docs/source/user-guide/latest/compatibility/floating-point.md @@ -23,22 +23,46 @@ Spark normalizes NaN and zero for floating point numbers for several cases. See However, one exception is comparison. Spark does not normalize NaN and zero when comparing values because they are handled well in Spark (e.g., `SQLOrderingUtil.compareFloats`). But the comparison functions of arrow-rs used by DataFusion do not normalize NaN and zero (e.g., [arrow::compute::kernels::cmp::eq](https://docs.rs/arrow/latest/arrow/compute/kernels/cmp/fn.eq.html#)). -For top-level `FLOAT` and `DOUBLE` comparisons, Comet normalizes both operands before native -execution, including noncanonical NaN literals. Top-level `IN`, `InSet`, and `NOT IN` membership -also normalize dynamic candidates and lists containing NaN. When every candidate is a non-NaN -literal, Comet keeps DataFusion's static filter and pruning path, enumerating both signed-zero -forms when a list contains zero. +For `FLOAT` and `DOUBLE` comparisons (`=`, `<>`, `<=>`, `<`, `<=`, `>` and `>=`), Comet +compares the operands in Spark's order natively, including noncanonical NaNs: `-0.0` equals +`0.0`, all NaNs are equal, and NaN sorts above every other value. This applies wherever a +comparison appears: projections, filters, aggregate arguments and `FILTER` clauses, join +conditions, sort keys, and generator arguments. + +A native Parquet scan prunes row groups and pages with its data filters. So that this pruning +still applies, a `FLOAT` or `DOUBLE` column compared with a constant other than NaN in a data +filter is compared without normalizing the column, and the filter above the scan evaluates the +comparison again with Spark's semantics. Every other comparison in a data filter follows Spark's semantics. +Bloom filters store `-0.0` and `0.0` separately, so `=` against either zero checks them for both. +With `spark.comet.parquet.rowFilterPushdown.enabled=true` the scan also drops the rows its data +filters reject, so every comparison in a data filter follows Spark's semantics, and a `FLOAT` or `DOUBLE` +comparison in a data filter does not prune row groups or pages +([#6702](https://github.com/apache/datafusion-comet/issues/6702)). + +Spark's own Parquet reader keeps `-0.0` and `0.0` apart when it prunes row groups with dictionaries +and bloom filters, so it can skip a row group holding rows that its filter matches. For example, +`WHERE d = -0.0D` can skip a row group that holds `0.0` but not `-0.0`, and so can `<=>` or `IN` +against `-0.0D`. `WHERE d >= 0.0D` can skip a dictionary-encoded row group whose largest value is +`-0.0`, and `WHERE d <= -0.0D` one whose smallest value is `0.0`. Comet reads those row groups, so +it can return rows that Spark does not. Comet's result is the one that Spark's comparison semantics +call for, and Spark returns the same rows with `spark.sql.parquet.filterPushdown=false`. + +Top-level `IN`, `InSet`, and `NOT IN` membership also normalize dynamic candidates and lists +containing NaN. When every candidate is a non-NaN literal, Comet keeps DataFusion's static filter +and pruning path, enumerating both signed-zero forms when a list contains zero. This scalar membership handling does not yet recurse into floating-point leaves nested in arrays or structs; see [#6019](https://github.com/apache/datafusion-comet/issues/6019). -## Nested equality and membership +## Nested comparisons and membership For arrays and structs containing `FLOAT` or `DOUBLE`, native `=`, `<>`, `IN`, and `NOT IN` compare signed zeros as equal and all NaN representations as equal, matching Spark. This also -covers single-candidate membership that Spark rewrites into equality. +covers single-candidate membership that Spark rewrites into equality. `<=>`, `<`, `<=`, `>`, and +`>=` follow Spark's order too, in which NaN sorts above every other value. -Equality and dynamic membership compare nested elements directly and stop at the first mismatch. +Comparisons and dynamic membership compare nested elements directly and stop at the first +mismatch. Constant membership sets use normalized comparison values for static lookup. These operations preserve SQL null semantics and do not change the values returned by projections. @@ -58,15 +82,14 @@ Because those comparison keys match Spark, `spark.comet.exec.strictFloatingPoint force a fallback for them: sort keys, window and rank order keys, and range partitioning keys all stay native under strict mode, whether the floats in them are scalar or nested. -The exception is a key that nests floats in an array or struct whose type can hold a null element -or field. Spark orders such a null below every other value, whatever the key's `NULLS FIRST` or -`NULLS LAST`. The native sort places it by that null order, so `ASC NULLS LAST` and -`DESC NULLS FIRST` can differ from Spark, and a `RANGE` window frame orders it above every other -value, so a running aggregate can span the whole partition -([#6476](https://github.com/apache/datafusion-comet/issues/6476), -[#6477](https://github.com/apache/datafusion-comet/issues/6477)). Strict mode makes those keys fall -back to Spark. A key whose type cannot hold a null, such as `array(coalesce(x, 0.0D))`, stays -native. +That includes a key that nests floats in an array or struct whose type can hold a null element or +field. Spark orders such a null below every other value, whatever the key's `NULLS FIRST` or +`NULLS LAST`, so Comet falls back in every mode, not only in strict mode, for the two shapes where +the native sort or window frame would place it differently: `ASC NULLS LAST` or `DESC NULLS FIRST` +on such a key ([#6476](https://github.com/apache/datafusion-comet/issues/6476)), and a `RANGE` +window frame that has to find a row's peers over it +([#6477](https://github.com/apache/datafusion-comet/issues/6477)). The +[operator compatibility notes](operators.md) describe both. `array_min` and `array_max` use Spark-compatible native comparisons in both strict and non-strict floating-point modes. Signed zeros compare equal, and all NaN representations compare equal and @@ -123,7 +146,12 @@ reports 4.2.0 but includes SPARK-59602 still runs natively; set `spark.comet.expression.ArrayDistinct.enabled=false` and `spark.comet.expression.ArrayUnion.enabled=false` on such a build. -## `array_remove` and `sort_array` +## `array_contains`, `array_remove` and `sort_array` + +`array_contains` compares `FLOAT` and `DOUBLE` elements as Spark does: `-0.0` equals `0.0`, and all +NaN representations are equal, inside nested arrays and structs too. The result keeps Spark's +three-valued form: null when nothing matches and the array holds a null element. Flat arrays run a +native kernel; nested float elements go through the codegen dispatcher. `array_remove` compares `FLOAT` and `DOUBLE` elements as Spark does: `-0.0` equals `0.0`, and all NaN representations are equal, inside nested arrays too. The elements it keeps retain their @@ -133,4 +161,4 @@ original NaN representations and zero signs. representations tie, and keeps equal elements in their original order. In an array whose elements can be null, such as one built from nullable columns, `-0.0` and `0.0` therefore tie. Sorting an array whose elements cannot be null in ascending order, Spark's generated code puts `-0.0` before -`0.0`, and so does Comet. Both expressions run natively in strict floating-point mode. +`0.0`, and so does Comet. All three expressions run natively in strict floating-point mode. diff --git a/docs/source/user-guide/latest/compatibility/operators.md b/docs/source/user-guide/latest/compatibility/operators.md index 4c65b2ee6c5..afab52bfce8 100644 --- a/docs/source/user-guide/latest/compatibility/operators.md +++ b/docs/source/user-guide/latest/compatibility/operators.md @@ -29,17 +29,20 @@ executed. Supported parent joins and aggregates remain eligible for native execution. Global aggregates still return one row (`COUNT = 0`, `SUM = NULL`), and grouped aggregates return no rows. Independent operator restrictions and aggregate buffer compatibility checks still apply. -Parquet writes whose input plans contain an empty relation use Spark's writer to preserve -readable empty output files and their schema metadata. + +A native Parquet write over a native empty relation stays native. Like Spark's writer, it runs one +task for the empty input, so the output still gets a schema-only Parquet file that readers can +infer the schema from. ## In-Memory Cache Comet can store cached relations (`df.cache()`, `CACHE TABLE`) in Arrow format and scan them -natively. This is experimental and disabled by default; see [In-Memory Cache](../in-memory-cache.md) -for how to enable it. Comet does not replace a `spark.sql.cache.serializer` that the application -has already set. Relations whose schema Comet's Arrow writer does not support are cached in -Spark's default format, and their scans fall back to Spark. Reads that feed Spark operators rather -than Comet operators can be slower than Spark's cache. +natively. This is enabled by default from Spark 3.5 but not on Spark 3.4; see +[In-Memory Cache](../in-memory-cache.md) for how to turn it off or on. Comet does not replace a +`spark.sql.cache.serializer` that the application has already set. Relations whose schema Comet's +Arrow writer does not support are cached in Spark's default format, and their scans fall back to +Spark. Reads that feed Spark operators rather than Comet operators can be slower than Spark's +cache. With Kryo and `spark.kryo.registrationRequired=true`, Comet needs its Kryo registrator whether or not the cache is enabled; see [Kryo serialization](../installation.md#kryo-serialization). diff --git a/docs/source/user-guide/latest/compatibility/spark-versions.md b/docs/source/user-guide/latest/compatibility/spark-versions.md index 337252cb794..84baf6194fa 100644 --- a/docs/source/user-guide/latest/compatibility/spark-versions.md +++ b/docs/source/user-guide/latest/compatibility/spark-versions.md @@ -89,11 +89,13 @@ Spark 4.2.0 is supported with Java 17 and Scala 2.13. - **`FILTER` on window aggregates**: Spark 4.2 accepts a `FILTER (WHERE ...)` clause on an aggregate function used as a window function. Comet does not support the clause there, so the window operator falls back to Spark. -- **`OneRowRelation` in `UNION` branches** - ([#4949](https://github.com/apache/datafusion-comet/issues/4949)): Spark 4.2 plans a - `OneRowRelation` into the `UNION` branches of some queries, for example TPC-DS q77a. Comet does - not convert `OneRowRelation` by default, so the union and the aggregates above it fall back to - Spark. +- **Merged subqueries in `UNION` branches** + ([#4949](https://github.com/apache/datafusion-comet/issues/4949)): Spark 4.2 turns a single-row + aggregate in some queries, for example TPC-DS q77a, into a scalar subquery that returns a struct, + and reads its fields in a projection over a `OneRowRelation`. Comet does not support a + struct-typed scalar subquery result yet + ([#5834](https://github.com/apache/datafusion-comet/issues/5834)), so that projection, and the + union and the aggregates above it, fall back to Spark. - **ANSI arithmetic overflow** ([#4967](https://github.com/apache/datafusion-comet/issues/4967)): Spark 4.2 changed some ANSI-mode arithmetic overflow behavior, and Comet does not match these changes yet. See diff --git a/docs/source/user-guide/latest/datasources.md b/docs/source/user-guide/latest/datasources.md index 30da3df12e1..f351aff5b6f 100644 --- a/docs/source/user-guide/latest/datasources.md +++ b/docs/source/user-guide/latest/datasources.md @@ -52,7 +52,8 @@ converted into Arrow format, allowing the Comet pipeline to take over after that ### Other Spark inputs Comet can also convert the output of these Spark inputs to Arrow format, so that the operators -above them can run in Comet. Each conversion is off by default. +above them can run in Comet. Only `spark.comet.convert.oneRowRelation.enabled` is on by default, +because the row it converts has no columns. - `spark.comet.convert.range.enabled`: `spark.range` and SQL `range()`, for ranges that Comet does not generate natively with `spark.comet.exec.range.enabled`. @@ -86,6 +87,11 @@ The same types apply to the output of typed `Dataset` operations, such as `map`, converts when `spark.comet.convert.typedDataset.enabled=true`. A column of any other type keeps the operators above the typed operation on Spark. +Comet does not convert a source, or the output of a typed `Dataset` operation, when the query uses +`input_file_name()`, `input_file_block_start()` or `input_file_block_length()`, so the operators +above it run in Spark. These functions report the file that Spark's reader is on when Spark evaluates +them, and the conversion and the Comet operators above it would read ahead of that. + ## Data Catalogs ### Apache Iceberg @@ -288,7 +294,10 @@ URLs. This is opt-in and disabled by default. Enable it by listing the schemes to treat as S3-compliant aliases in `spark.hadoop.fs.comet.s3Compliant.schemes` (Hadoop key `fs.comet.s3Compliant.schemes`), a comma-separated, case-insensitive list. This mirrors the -existing `fs.comet.libhdfs.schemes` config. +existing `fs.comet.libhdfs.schemes` config. The list entries are case-insensitive, but an Iceberg +table location must write the alias in lowercase (`blob://`, not `BLOB://`): the native Iceberg +reader opens each recorded location as written, and its S3 backend accepts only a lowercase +scheme prefix, so Comet declines such a location up front and leaves that scan to Spark. ```shell --conf spark.hadoop.fs.comet.s3Compliant.schemes=blob diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index 74232aac652..08b355ee7b0 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -155,8 +155,8 @@ The tables below list every Spark built-in expression with its current status. | `array` | ✅ | Native | | | `array_append` | ✅ | Native | | | `array_compact` | ✅ | — | | -| `array_contains` | ✅ | Native | Float/double element arrays route through the JVM codegen dispatcher by default; the native path is opt-in via allowIncompatible | -| `array_distinct` | ✅ | Native | Floating-point elements fall back on Spark versions other than 4.2.0; signed-zero and NaN results may differ with native opt-in ([details](compatibility/floating-point.md)) | +| `array_contains` | ✅ | Hybrid | Flat float/double element arrays use a native kernel with Spark's equality; nested float elements, non-default string collations, and a value that is not a literal or column over a nullable array, route through the JVM codegen dispatcher | +| `array_distinct` | ✅ | Native | Floating-point elements fall back on Spark versions other than 4.2.0; signed-zero and NaN results may differ with native opt-in ([details](compatibility/floating-point.md)). Non-default string collations fall back | | `array_except` | ✅ | Hybrid | Routes through the JVM codegen dispatcher by default; the incompatible native path is opt-in via allowIncompatible ([details](compatibility/expressions/array.md)) | | `array_insert` | ✅ | Native | | | `array_intersect` | ✅ | Hybrid | Routes through the JVM codegen dispatcher by default; the incompatible native path is opt-in via allowIncompatible ([details](compatibility/expressions/array.md)) | @@ -167,8 +167,8 @@ The tables below list every Spark built-in expression with its current status. | `array_prepend` | ✅ | — | | | `array_remove` | ✅ | Native | | | `array_repeat` | ✅ | Native | | -| `array_union` | ✅ | Native | Floating-point elements fall back on Spark versions other than 4.2.0; signed-zero and NaN results may differ with native opt-in ([details](compatibility/floating-point.md)) | -| `arrays_overlap` | ✅ | Native | | +| `array_union` | ✅ | Native | Floating-point elements fall back on Spark versions other than 4.2.0; signed-zero and NaN results may differ with native opt-in ([details](compatibility/floating-point.md)). Non-default string collations fall back | +| `arrays_overlap` | ✅ | Hybrid | Non-default string collations use the JVM codegen dispatcher | | `arrays_zip` | ✅ | Native | | | `element_at` | ✅ | Native | | | `flatten` | ✅ | Native | Binary/struct/map elements fall back | @@ -336,8 +336,7 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci ## generator_funcs `explode`, `explode_outer`, `posexplode`, and `posexplode_outer` are supported via -`CometExplodeExec` (operator-level, not expression-level) for array input; map input falls back -to Spark ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). Enabled by default via +`CometExplodeExec` for array and map inputs (operator-level, not expression-level). Enabled by default via `spark.comet.exec.explode.enabled`. | Function | Status | Implementation | Notes | diff --git a/docs/source/user-guide/latest/iceberg-writes.md b/docs/source/user-guide/latest/iceberg-writes.md index a89c397d493..a04f6be61cf 100644 --- a/docs/source/user-guide/latest/iceberg-writes.md +++ b/docs/source/user-guide/latest/iceberg-writes.md @@ -119,10 +119,15 @@ For an unpartitioned copy-on-write `MERGE`, the native Iceberg writer is reachab On Spark 4.1+, the native MergeRows path also preserves the semantic counters required by the summary-aware writer commit contract. -The mechanism behind row-level DML differs by Spark version: on Spark 4.0+ the analyzer emits -operation-coded rows that Comet's writer dispatches through `ReplaceData`'s projections, while -on Spark 3.4/3.5 the rewritten rows are written as a plain row stream. The supported set of -operations is the same either way. +For copy-on-write row-level DML, the mechanism differs by Spark version: on Spark 4.0+ the +analyzer emits operation-coded rows that Comet's writer dispatches through `ReplaceData`'s +projections, while on Spark 3.4/3.5 the rewritten rows are written as a plain row stream. The +supported set of operations is the same either way. + +On Spark 3.5+, merge-on-read uses Spark's `WriteDelta`. The split plan intercepts that command so +Comet can keep the same driver commit and reporting path, but task-side row-level writes stay on +Iceberg's JVM `DeltaWriter`; `CometIcebergWriteExec` is never used for position-delta rows. +Spark 3.4 leaves `WriteDelta` on Spark's stock write plan. On Spark 4.1+ the split plan matches two further stock-Spark behaviours: MERGE metrics are forwarded to the writer's commit (Iceberg 1.11+ records them in the snapshot summary), and @@ -134,8 +139,9 @@ changes. The rewrite is skipped — and the write runs through Spark's stock combined operator — when: - `spark.comet.write.iceberg.splitOperator.enabled` is `false` (the default); -- the write is not an Iceberg `SparkWrite` (any other V2 data source); -- the table uses merge-on-read: delta writes (Iceberg `WriteDelta`) are not intercepted; +- the write is neither an Iceberg `SparkWrite` nor a supported Iceberg position-delta write; +- the table uses merge-on-read on Spark 3.4; Spark 3.5+ `WriteDelta` is intercepted but remains + on Iceberg's JVM `DeltaWriter`; - the statement is CTAS / RTAS on Spark 3.4, where the staged exec writes inline; on Spark 3.5+ those statements re-plan their inner append, which is intercepted normally; - the write requires Spark's commit coordinator, which Comet's per-task commit protocol does @@ -273,8 +279,9 @@ Other `write.*` properties are intentionally not gated because they cannot make writer produce different data files: distribution and ordering settings shape the Spark plan identically on both paths, WAP / branch / snapshot properties act on the JVM committer, `write.avro.*` / `write.orc.*` apply only to formats already excluded, and merge-on-read -settings route the write through `WriteDelta`, which the split plan never intercepts. Every -rule is pinned by `CometIcebergWriteDetectionSuite`. +settings route the write through `WriteDelta`, whose task writer remains iceberg-java rather +than the native writer. Every native eligibility rule is pinned by +`CometIcebergWriteDetectionSuite`. Manifest `DataFile` metrics are assembled on the JVM before commit: each written file's metrics are re-derived from its parquet footer through the version-matched @@ -299,6 +306,18 @@ failure), the error propagates as an ordinary Spark task failure and Spark's tas re-executes it — through the native writer again. Retries cannot collide: each attempt's task attempt id is embedded in its data file names. +The native writer's buffers are charged to Comet's memory pool, the off-heap budget Comet's other +native operators draw on, where iceberg-java's buffers sit on the JVM heap. A fanout write keeps a +data file open for every partition a task writes to. Each open file holds the row group it is +writing in memory, up to `write.parquet.row-group-size-bytes`, and on S3 or GCS also the last row +group it flushed, which is uploaded once the next one is complete or the file closes. So a task +writing to many partitions needs memory in proportion to them. When the pool cannot grant it, the +task fails with a `CometNativeException` reading `Additional allocation failed for IcebergWriteExec` +instead of exceeding the executor's memory, and Spark retries it like any other task failure. Such a +write fits in less memory with the fanout writer disabled (`write.spark.fanout.enabled=false`): +Spark then sorts each task's rows by partition, and the task keeps one file open at a time. A +smaller row-group size also helps. Otherwise the write needs a larger `spark.memory.offHeap.size`. + Partial results are never committed. The commit set is exactly the commit messages returned by successful tasks — a failed task contributes none — and if the job fails, the driver-side commit operator aborts without committing anything. A failed task attempt also deletes the diff --git a/docs/source/user-guide/latest/in-memory-cache.md b/docs/source/user-guide/latest/in-memory-cache.md index 20091ac3172..d6867defeab 100644 --- a/docs/source/user-guide/latest/in-memory-cache.md +++ b/docs/source/user-guide/latest/in-memory-cache.md @@ -21,18 +21,24 @@ Comet can store Spark's in-memory cache (`CACHE TABLE`, `df.cache()`, `df.persist()`) in an Arrow format that Comet operators read directly. Without it, a cached table is stored in Spark's own -format and every scan of it has to convert each batch before Comet can continue, which shows up in -the plan as a `CometSparkColumnarToColumnar` above the cache scan. +format, which Comet operators cannot read. Under Comet's default settings the operators above the +cache scan then run on Spark. With `spark.comet.convert.inMemoryCache.enabled`, a +`CometSparkColumnarToColumnar` above the scan converts each batch for Comet operators instead. -This feature is **experimental and disabled by default**. Turn it on at startup, alongside the rest -of Comet's configuration: +This feature is **enabled by default** from Spark 3.5. To turn it off, set the config at startup, +alongside the rest of Comet's configuration: ```shell $SPARK_HOME/bin/spark-shell \ ... \ - --conf spark.comet.exec.inMemoryCache.enabled=true + --conf spark.comet.exec.inMemoryCache.enabled=false ``` +On Spark 3.4 it is disabled by default, and setting the config to `true` turns it on. Spark 3.4 has +no hook for the rule through which Comet lets AQE coalesce the shuffle partitions of a union that +Comet runs, so a union of a shuffle and a relation cached in Comet's format keeps every shuffle +partition there. + It has to be set before the `SparkContext` starts. Comet's driver plugin chooses `spark.sql.cache.serializer` once, while the context is initializing, so an application keeps the cache format it started with however the config is set afterwards. The plugin installs Comet's @@ -112,7 +118,7 @@ nowhere to record either that a column is dictionary encoded or the dictionary i | Config | Default | Description | | ------------------------------------------------------- | ------- | ---------------------------------------------------------------------------------------------------------------------------------------------- | -| `spark.comet.exec.inMemoryCache.enabled` | `false` | Whether to store and scan Spark's in-memory cache in Comet's format. Read at startup. | +| `spark.comet.exec.inMemoryCache.enabled` | `true` | Whether to store and scan Spark's in-memory cache in Comet's format. Read at startup. Defaults to `false` on Spark 3.4. | | `spark.comet.exec.inMemoryCache.compression.codec` | `zstd` | Arrow IPC compression codec for cached data: `zstd` or `none`. Affects newly cached data only — a batch records the codec it was written with. | | `spark.comet.exec.inMemoryCache.compression.zstd.level` | `1` | Compression level when the codec is `zstd`. Ignored otherwise. | @@ -161,10 +167,57 @@ back to Spark row execution above the scan and the two columns stop measuring th Read what this compares carefully. Comet execution is on in both columns, so the aggregation runs on Comet either way and only the cache-scan boundary moves: on the left, Spark's `InMemoryTableScanExec` feeds those same Comet operators through a `CometSparkColumnarToColumnar` -bridge; on the right, `CometInMemoryTableScan` feeds them directly. Both columns read the same +bridge, which the benchmark turns on with `spark.comet.convert.inMemoryCache.enabled`; on the +right, `CometInMemoryTableScan` feeds them directly. Both columns read the same Comet-written `CometCachedBatch`. These numbers are therefore "keep the cached scan native" against "fall back to a Spark cache scan and convert", not Comet against Spark execution, and not a -comparison with Spark's own cache format. That comparison is under [Limitations](#limitations). +comparison with Spark's own cache format, which follows. + +### Against Spark's cache format + +What turning the feature on changes for a query that Comet runs is measured against Spark's own +cache format by the benchmark's adaptive cases. Comet and AQE are on, Comet's other settings are at +their defaults, and the same 5M-row relation is cached in each format. The defaults leave +`spark.comet.convert.inMemoryCache.enabled` off, so Comet operators cannot read Spark's cache scan, +and with Spark's format the operators directly above the scan run on Spark. Measured on an AMD +Ryzen 9 7950X3D (JDK 17, Spark 4.1, release build): + +| Query shape | Spark's cache format | Comet's cache format | Relative | +| -------------------------- | -------------------: | -------------------: | -------: | +| Row count only (0 of 6) | 29 ms | 24 ms | 1.2x | +| Narrow projection (1 of 6) | 52 ms | 34 ms | 1.5x | +| 3 of 6 columns | 102 ms | 112 ms | 0.9x | +| Full projection (6 of 6) | 299 ms | 224 ms | 1.3x | + +A Spark operator above the cache scan, standing in for any operator Comet does not support, is +measured the same way, with Comet's aggregate turned off. With Comet's format, the native scan feeds +that operator through a columnar-to-row transition: + +| Query shape | Spark's cache format | Comet's cache format | Relative | +| -------------------------- | -------------------: | -------------------: | -------: | +| Row count only (0 of 6) | 39 ms | 16 ms | 2.4x | +| Narrow projection (1 of 6) | 50 ms | 27 ms | 1.8x | +| 3 of 6 columns | 97 ms | 112 ms | 0.9x | +| Full projection (6 of 6) | 303 ms | 299 ms | 1.0x | + +Comet's format is as fast or faster in every shape but one: the read of three of the six columns, +all of them longs, is about 10% slower under either kind of operator. That cost is `zstd` +decompression. With the `none` codec, the same read is 2.7x faster than Spark's format with Comet +operators above the scan, and 1.6x faster with a Spark operator above it. + +Building the cache is measured the same way, from the same source, with its rows produced either by +Comet operators or by Spark operators. Comet's format is written straight from the Arrow batches of +Comet operators, while Spark's is always built from rows, so above Comet operators it first converts +their batches. Measured on an Apple M3 Max (JDK 17, Spark 4.1, release build): + +| Rows produced by | Spark's cache format | Comet's cache format, `zstd` | Comet's cache format, `none` | +| ---------------- | -------------------: | ---------------------------: | ---------------------------: | +| Comet operators | 3285 ms | 944 ms | 601 ms | +| Spark operators | 3680 ms | 2079 ms | 1710 ms | + +Spark's format holds the relation in 217 MiB, and Comet's in 51 to 55 MiB with `zstd` and 315 MiB +with `none`. So with the default codec, Comet's format builds 1.8x to 3.5x faster than Spark's and +takes about a quarter of the memory. ## Kryo @@ -220,6 +273,9 @@ operator reads a relation cached in Comet's format depends on the scan below it: and operators that do not take part in code generation, such as exchanges and limits, or a query that returns the cached rows as they are. +Whenever Spark's scan reads Comet's format, which also happens for a relation whose cached plan +records `Dataset.observe` metrics, Comet records a fallback reason on the scan. + Measured by the same benchmark over the same 5M-row relation, with native execution off so that Spark operators consume the cached data, Comet disabled for the row reader and enabled for the fused reader (Apple M4, JDK 17, Spark 4.1; the average of two runs): @@ -234,9 +290,9 @@ reader (Apple M4, JDK 17, Spark 4.1; the average of two runs): The fused reader is faster than Spark's own format for the narrowest reads and within 20% of it for the others. The row reader takes up to 1.6 times as long, and it is the only reader for relations wider than `spark.sql.codegen.maxFields`: reading every column of relations of 100, 200 and 1500 -nullable `bigint` columns took 2.2 to 2.5 times as long as from Spark's format. These gaps are why -the feature is still off by default; -[#5485](https://github.com/apache/datafusion-comet/issues/5485) tracks them. +nullable `bigint` columns took 2.2 to 2.5 times as long as from Spark's format. A Spark operator +above Comet's native cache scan does not pay this; see [Performance](#performance). +[#5485](https://github.com/apache/datafusion-comet/issues/5485) tracks these gaps. Comet's serializer exists because Spark's own Arrow cache format ([SPARK-57268](https://issues.apache.org/jira/browse/SPARK-57268)) is only available from Spark diff --git a/docs/source/user-guide/latest/migration-guide.md b/docs/source/user-guide/latest/migration-guide.md index b55cd94d97c..c7ceea0ee60 100644 --- a/docs/source/user-guide/latest/migration-guide.md +++ b/docs/source/user-guide/latest/migration-guide.md @@ -59,6 +59,34 @@ key is removed. ## Upgrading to Comet 1.2.0 +Comet `1.2.0` makes no behavior changes that need a `spark.comet.legacy.*` key. The changes below +need none either, but check whether any of them applies to your deployment. + +### In-Memory Cache Enabled by Default + +`spark.comet.exec.inMemoryCache.enabled` now defaults to `true` on Spark 3.5 and later. There, an +application that loads `CometPlugin` now stores what it caches with `CACHE TABLE`, `df.cache()` or +`df.persist()` in Comet's Arrow format instead of Spark's, and Comet scans it natively. On Spark 3.4 +the default stays `false`; see [In-Memory Cache](in-memory-cache.md) for why. The format does not +change query results, but it can change performance. Spark's own cache scan can read Comet's format +more slowly than Spark's, which matters when a session turns Comet or its native execution off +after caching, and Comet records a fallback reason on such a scan. See +[In-Memory Cache](in-memory-cache.md#limitations). + +The format is chosen once, when the application starts. To keep Spark's format, set +`spark.comet.exec.inMemoryCache.enabled=false` then. Comet also keeps Spark's format without that +setting when the application: + +- starts with `spark.comet.enabled` or `spark.comet.exec.enabled` set to `false`. +- leaves Comet shuffle enabled without one of Comet's shuffle managers, so that Comet disables + itself. +- uses Kryo with `spark.kryo.registrationRequired=true` and has not registered Comet's cached + batch, because Kryo would reject it. To use Comet's format, register Comet's classes with + `spark.kryo.registrator=org.apache.comet.CometKryoRegistrator`; see + [Kryo](in-memory-cache.md#kryo). + +An application that sets `spark.sql.cache.serializer` itself keeps the serializer it chose. + ### Deprecated and Removed Settings Using `spark.comet.sparkToColumnar.enabled` and `spark.comet.sparkToColumnar.supportedOperatorList` @@ -79,8 +107,9 @@ set, `spark.comet.sparkToColumnar.enabled=true` still converts the first four, w by default, and a list that names any of the five still converts it, as before. In both cases the driver logs a warning that names the config to use instead. Before Spark 4.1, Spark plans a query without a `FROM` clause as an `RDDScan`, so on those versions `RDDScan` in the list also converted -it. The list remains the way to convert other leaf operators, such as the scan of a Data Source V2 -connector. +it. `spark.comet.convert.oneRowRelation.enabled` is on by default, so such a query is now converted +without either setting, and logs no warning. The list remains the way to convert other leaf +operators, such as the scan of a Data Source V2 connector. ## Upgrading to Comet 1.1.0 diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index 0155ebb8dae..fbe961527f8 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -64,9 +64,9 @@ omitted from the tables below and may be reconsidered based on demand: | `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. | +| `EmptyRelationExec` | ✅ | Spark 4.0 and later. See [Empty Relations](compatibility/operators.md#empty-relations) for native-input support, including native Parquet writes. | | `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). | +| `InMemoryTableScanExec` | ✅ | Enabled by default from Spark 3.5. Comet installs its Arrow cache serializer as the application starts if `spark.comet.exec.inMemoryCache.enabled` is true, its default except on Spark 3.4. 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 @@ -123,12 +123,12 @@ natively on `BroadcastHashJoinExec` and `ShuffledHashJoinExec`. Existence sort-m ## Generators and set operations -| Operator | Status | Notes | -| -------------- | ------ | ---------------------------------------------------------------------------------------------------------------- | -| `GenerateExec` | ✅ | Supports `explode`, `explode_outer`, `posexplode`, `posexplode_outer` over arrays. `inline` / `stack` fall back. | -| `ExpandExec` | ✅ | | -| `UnionExec` | ✅ | | -| `CoalesceExec` | ✅ | | +| Operator | Status | Notes | +| -------------- | ------ | ------------------------------------------------------------------------------------------------------------------------- | +| `GenerateExec` | ✅ | Supports `explode`, `explode_outer`, `posexplode`, `posexplode_outer` over arrays and maps. `inline` / `stack` fall back. | +| `ExpandExec` | ✅ | | +| `UnionExec` | ✅ | | +| `CoalesceExec` | ✅ | | ## Writes diff --git a/docs/source/user-guide/latest/rust_udfs.md b/docs/source/user-guide/latest/rust_udfs.md index ba98cd6acbb..147a2069c94 100644 --- a/docs/source/user-guide/latest/rust_udfs.md +++ b/docs/source/user-guide/latest/rust_udfs.md @@ -158,6 +158,7 @@ CometNativeUDF.register( Registration loads the library on the driver and verifies that a function with that name exists, so a bad path or a missing function fails immediately with a clear error rather than at execution time. +It also checks that Comet can carry the argument and return types natively. The function is then callable from SQL or the DataFrame API like any other: @@ -181,6 +182,19 @@ way it would for any other shared object, through `LD_LIBRARY_PATH` on Linux and paths may be loaded, so it makes no difference to the trust decision described under [Limitations](#limitations). +### When Spark evaluates the call + +A call runs only inside Comet's native execution. Spark cannot run the function, so wherever Spark +has to evaluate the call itself, the query fails with `CometUdfNotEvaluatedException`. That happens: + +- when Comet does not take the operator holding the call, for example because another expression + in it is not supported, or because the call is an argument of an ordinary Scala or Java UDF. The + query's extended explain output gives the reason. +- while Spark plans the query: over local data such as `VALUES`, which the optimizer evaluates + eagerly, in a filter on partition columns, which Spark evaluates to prune partitions, and in the + sort keys of a global sort, which Spark evaluates on a sample of rows to choose range bounds. To + sort on a UDF's result, select it as a column and sort on that column instead. + ## Return types A UDF does not have one fixed return type. Its `return_field` is called with the actual argument @@ -277,6 +291,9 @@ the first panic, and the SDK cannot prevent it. This feature is at an early stage. The current limitations are: - **Scalar functions only.** Aggregate, window, and table functions are not supported. +- **Spark cannot evaluate a call.** A query fails wherever Spark would have to run the function + itself, which includes a few places while it plans the query. See + [When Spark evaluates the call](#when-spark-evaluates-the-call). - **Immutable functions only.** A UDF must return the same output for the same input. Comet plans every Rust UDF with DataFusion's `Volatility::Immutable`, so a function that reads a clock, draws from an RNG, or carries state across batches may be folded at plan time, evaluated once and diff --git a/docs/source/user-guide/latest/tuning.md b/docs/source/user-guide/latest/tuning.md index 158883bb6c1..28f6ce112a7 100644 --- a/docs/source/user-guide/latest/tuning.md +++ b/docs/source/user-guide/latest/tuning.md @@ -36,6 +36,21 @@ guide is split into the following pages: - [Reducing Row/Columnar Conversion Overhead](tuning/transitions.md): stages in which many operators fall back to Spark. +## Nested and Wide Data + +Rows with many columns, long strings, or nested columns such as arrays of structs make every batch +larger, so queries over such data spill more. For these workloads: + +- Choose the codec for spill files from native sorts, aggregations, and joins, which Comet + compresses with `lz4` by default. See + [Compressing Spill Files](tuning/memory.md#compressing-spill-files). +- Lower the batch size if wide batches cause spilling or out-of-memory errors. See + [Batch Size](tuning/memory.md#batch-size). +- Check which shuffles can run natively when partitioning keys are nested. See + [Native Shuffle](tuning/shuffle.md#native-shuffle). +- Watch for row/columnar conversions, which are expensive for such schemas. See + [Wide or Deeply Nested Schemas](tuning/transitions.md#wide-or-deeply-nested-schemas). + ## Configuring Tokio Runtime Comet uses a global tokio runtime per executor process. By default it starts one worker thread per executor core diff --git a/docs/source/user-guide/latest/tuning/memory.md b/docs/source/user-guide/latest/tuning/memory.md index 1f6ea98a2be..2632369df24 100644 --- a/docs/source/user-guide/latest/tuning/memory.md +++ b/docs/source/user-guide/latest/tuning/memory.md @@ -282,3 +282,13 @@ use more than `N` times this value on shared local disks. If the limit is reache fail and the query errors out. Raise this on workloads with large sort/aggregate/shuffle spills, or lower it to protect executors on shared disks, remembering that the total across an executor is a multiple of this value. + +## Compressing Spill Files + +`spark.comet.exec.spill.compression.codec` sets the codec for the files that native sorts, +aggregations, sort-merge joins, and nested loop joins spill to disk. The default is `lz4`, which +Spark also uses for its own spill files by default. `zstd` writes smaller files, which reduces +spill I/O and how much of `spark.comet.maxTempDirectorySize` they use, but takes more CPU to +compress and decompress them. `none` writes them uncompressed. Comet's shuffle writers use +`spark.comet.shuffle.compression.codec` instead. See +[Shuffle Compression](shuffle.md#shuffle-compression). diff --git a/docs/source/user-guide/latest/tuning/operators.md b/docs/source/user-guide/latest/tuning/operators.md index 644f35483ba..cb538fb63d4 100644 --- a/docs/source/user-guide/latest/tuning/operators.md +++ b/docs/source/user-guide/latest/tuning/operators.md @@ -30,6 +30,14 @@ to configure Comet to convert `SortMergeJoin` to `ShuffledHashJoin`. Comet does to test with both for your specific workloads. To configure Comet to convert `SortMergeJoin` to `ShuffledHashJoin`, set `spark.comet.exec.forceShuffledHashJoin=true`. +The conversion only happens when the build side is under a size limit, and a join whose build side has no statistics +is left as a `SortMergeJoin`. The size is Spark's planning estimate of the build side, or under AQE the materialized +shuffle size of the build side. By default the limit is Spark's own rule for choosing a `ShuffledHashJoin`: +`spark.sql.autoBroadcastJoinThreshold` times the initial shuffle partition count +(`spark.sql.adaptive.coalescePartitions.initialPartitionNum` when AQE and partition coalescing are both on and it is +set, else `spark.sql.shuffle.partitions`). When broadcasts are disabled with a non-positive threshold, Spark's default +threshold of 10 MB is used instead. Set `spark.comet.exec.forceShuffledHashJoin.maxBuildSize` to a size in bytes to use a fixed +limit, or to a non-positive value to convert every eligible join regardless of size. ### Join Runtime Filters @@ -154,12 +162,12 @@ nested in arrays and structs, so `ORDER BY`, window ordering and range partition even with `spark.comet.exec.strictFloatingPoint=true`. Only the comparison key is normalized; returned values keep their original NaN representation and zero sign. -The exception is a key that nests floating-point values in an array or struct whose type can hold a null element or -field. Spark orders such a null below every other value, and the native sort and `RANGE` window frames do not +A key that nests values in an array or struct whose type can hold a null element or field falls back in every mode under +`ASC NULLS LAST` or `DESC NULLS FIRST`, and in a `RANGE` window frame that has to find a row's peers, because Spark orders +such a null below every other value and the native sort and window frames would not ([#6476](https://github.com/apache/datafusion-comet/issues/6476), -[#6477](https://github.com/apache/datafusion-comet/issues/6477)), so -`spark.comet.exec.strictFloatingPoint=true` makes those keys fall back to Spark. They can be forced back onto the -native path with `spark.comet.expression.SortOrder.allowIncompatible=true`. +[#6477](https://github.com/apache/datafusion-comet/issues/6477)). The default null orders run natively, and +`spark.comet.expression.SortOrder.allowIncompatible=true` forces the other null orders onto the native path. `sort_array` sorts array elements rather than ordering rows. It follows Spark's floating-point ordering as well, so it also stays native with `spark.comet.exec.strictFloatingPoint=true`. diff --git a/docs/source/user-guide/latest/tuning/scans.md b/docs/source/user-guide/latest/tuning/scans.md index f23d32f41c7..40bc82634ec 100644 --- a/docs/source/user-guide/latest/tuning/scans.md +++ b/docs/source/user-guide/latest/tuning/scans.md @@ -28,7 +28,8 @@ decode step and lazily materializes projected columns for surviving rows. This c CPU and memory when the filter is highly selective on a small subset of columns. It is disabled by default because it can hurt when the filter is not selective or when most columns must be read anyway. Row-group, page-index, and bloom-filter pruning happen regardless of this flag whenever Spark's -`spark.sql.parquet.filterPushdown` is on. +`spark.sql.parquet.filterPushdown` is on, except for filters that compare `FLOAT` or `DOUBLE` values, +which do not prune while this flag is on (see [Floating-point Number Comparison](../compatibility/floating-point.md)). ### Parquet Native Scans diff --git a/docs/source/user-guide/latest/tuning/shuffle.md b/docs/source/user-guide/latest/tuning/shuffle.md index 71c9c431ab0..b7ad3d2f090 100644 --- a/docs/source/user-guide/latest/tuning/shuffle.md +++ b/docs/source/user-guide/latest/tuning/shuffle.md @@ -66,23 +66,34 @@ write side. To restrict columnar shuffle to cases where the child is already a C as native Spark shuffles, which avoids the row to columnar conversion but means the downstream stage will also start on Spark. +Setting `spark.comet.convert.shuffleInput.enabled=true` makes such a shuffle use native shuffle instead. Comet +converts the child's rows to Arrow, as the `spark.comet.convert.*` settings do for Spark scans, and shuffles the +batches natively. With hash partitioning that is usually faster than columnar shuffle, while with range partitioning +the two come out about even. This applies only where native shuffle supports the partitioning and Comet can convert +all of the child's columns, which rules out calendar intervals, and arrays and maps other than `array` and +`map`. A shuffle that hashes a decimal wider than 18 digits, or a string or a value computed from a +string, stays on columnar shuffle, so that it partitions rows as Spark does. The converted batches are allocated on +the JVM outside Comet's memory pools. The setting is disabled by default. + ### Automatic Revert to Spark Shuffle -When a Comet columnar shuffle ends up between a partial and a final aggregate that Comet could not convert (both -remain Spark `HashAggregateExec` or `ObjectHashAggregateExec` operators), Comet reverts it to Spark's built-in shuffle. +When a Comet columnar shuffle, or a native shuffle over rows that `spark.comet.convert.shuffleInput.enabled` converted, +ends up between a partial and a final aggregate that Comet could not convert (both remain Spark `HashAggregateExec`, +`ObjectHashAggregateExec` or `SortAggregateExec` operators), Comet reverts it to Spark's built-in shuffle. Keeping columnar shuffle between the two row-based aggregates would add `row -> Arrow -> shuffle -> Arrow -> row` -conversions with no Comet consumer on either side to benefit from columnar output. Other shuffles between non-Comet -operators are not reverted. +conversions with no Comet consumer on either side to benefit from columnar output. The final aggregate has to read the +shuffle directly, which a `SortAggregateExec` with grouping keys does not: it reads the shuffle through a sort, so that +shuffle is not reverted. Other shuffles between non-Comet operators are not reverted either. This shifts the affected shuffles from Comet's off-heap memory pool back to the JVM execution memory pool. Clusters tuned for a small JVM heap may see `ExternalSorter` spills on queries where this revert fires. Shuffle I/O may also grow marginally because Spark's row-based serializer generally compresses less well than Comet's Arrow IPC format. -Each revert is logged at `INFO` level on the driver as `Reverting Comet columnar shuffle to Spark shuffle between - and `, which lets you correlate any unexpected behavior with this optimization. +Each revert is logged at `INFO` level on the driver as `Reverting Comet shuffle to Spark shuffle between +and `, which lets you correlate any unexpected behavior with this optimization. This optimization is enabled by default and can be disabled by setting -`spark.comet.shuffle.revertRedundantColumnar.enabled=false`, in which case Comet will keep the columnar shuffle +`spark.comet.shuffle.revertRedundantColumnar.enabled=false`, in which case Comet will keep its shuffle even when both of those aggregates run on Spark. ## Shuffle Compression diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index 13733f0b510..4791f1abf5a 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -298,7 +298,7 @@ by role. Names match what is shown in the plan output. | `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). | +| `CometInMemoryTableScan` | JVM-side scan of a table cached in Comet's Arrow format (enabled by default from Spark 3.5). | | `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). | diff --git a/docs/source/user-guide/older-versions.md b/docs/source/user-guide/older-versions.md index 00d936be54c..ca46ba134ec 100644 --- a/docs/source/user-guide/older-versions.md +++ b/docs/source/user-guide/older-versions.md @@ -20,11 +20,12 @@ under the License. # Older Versions These user guides are kept for reference. They are not maintained: use the -[current stable release](1.0/index) for production guidance. +[current stable release](1.1/index) for production guidance. ```{toctree} :maxdepth: 1 +1.0.x <1.0/index> 0.17.x <0.17/index> 0.16.x <0.16/index> 0.15.x <0.15/index> diff --git a/native/Cargo.lock b/native/Cargo.lock index e777e80d00c..e4d2aac7e1d 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -1881,9 +1881,9 @@ dependencies = [ [[package]] name = "datafusion" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "14313430439ef858ed5f52ca9c840e9eea9844da34d3716d05a476449634da5d" +checksum = "4cd5317db77ba6bca25c761a8091b9d294a9f4188e4f9dcd0243e6a57de862ac" dependencies = [ "arrow", "arrow-schema", @@ -1930,9 +1930,9 @@ dependencies = [ [[package]] name = "datafusion-catalog" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2c0756ebc770bcebae366dfd273dad50935f6103b4d212253cb17e2e53c5427f" +checksum = "f35d36794509b7bfcbf90e11a1b5e4d047123c34610105e2e80a2048351cdb2f" dependencies = [ "arrow", "async-trait", @@ -1955,9 +1955,9 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bbde2961698705d8402f0f382f08876fc6f1126ffde790ef17833cf0ab87cd59" +checksum = "a6454bfb7f43b83942e2ca2576fd09342eedbd0d323e1eef72c3c8e4d57d7203" dependencies = [ "arrow", "async-trait", @@ -2143,6 +2143,7 @@ dependencies = [ "datafusion", "datafusion-comet-common", "datafusion-comet-jni-bridge", + "datafusion-functions", "datafusion-spark", "datasketches", "futures", @@ -2163,9 +2164,9 @@ dependencies = [ [[package]] name = "datafusion-common" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a6e06fe711ff815c8cf3143158547aea599702fd43f76f5866c70d5a5b3693a0" +checksum = "d26626e12b224b83351239465dcd00f72d47b02fb5e5b65a643e2dc8ed37f687" dependencies = [ "arrow", "arrow-ipc", @@ -2190,9 +2191,9 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f7a5107acff4e3ed50a7732acec06bba40e74ab49df021397ed250d60f04e42a" +checksum = "a106420cbfa980f242c22243ae9a6cce661d4d86ce96cff87dacabf7754a4470" dependencies = [ "futures", "log", @@ -2201,9 +2202,9 @@ dependencies = [ [[package]] name = "datafusion-datasource" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f69f45d2d700785be8df9d4b44410e1c8e4a856f10cc728b418a1066c0b594fb" +checksum = "210aa1e9525e9454f2399dbb8b0b95510145bd8c1fdc35de74e9ccbf81bcb699" dependencies = [ "arrow", "async-compression", @@ -2237,9 +2238,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "81650af64894ef34b46d2ef818ef244b40a81d90818f3bc642bf8fbf6758f8e0" +checksum = "65b1b041688cad707fa103c3e7f4c6d89b3134ebf67da0129070dad3381851eb" dependencies = [ "arrow", "arrow-ipc", @@ -2261,9 +2262,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5e2519bed59a6907c79298b9c21fe807bbc34e4ff8414577d699b0741083593" +checksum = "2a7b34e253c4fa5251f174977af910a2f2ec831ec39dbbb56b29ce505cff278b" dependencies = [ "arrow", "async-trait", @@ -2284,9 +2285,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "635e61290e171597f6e50427e0cb486eda468ab47e880fa24d15bbb0fa24ac7e" +checksum = "1663338f43be6a5746473008e2d0fbde282577ddf0bae90d6c7386e988d71b29" dependencies = [ "arrow", "async-trait", @@ -2307,9 +2308,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07cd518f85b3a027d913047b4ab0811111a2b734fdf84a51b78234455f10d8fa" +checksum = "7c521ae91b6d899d7190812881f45fe9463c222c11d90149f5332a2e56555904" dependencies = [ "arrow", "arrow-schema", @@ -2339,15 +2340,15 @@ dependencies = [ [[package]] name = "datafusion-doc" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d44c5919b560862dd1cb20a5c75b3e48c3d9e49719c91d78da5c38eaad5495c4" +checksum = "0c2569c2e74a34fcae3a369ab3912d61e41a062592510f11e31b1dae2725ecdc" [[package]] name = "datafusion-execution" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7620ad535af8730186bdec05d720e40d9a4584b40c7eb75f72c43aa7978e5484" +checksum = "39cfdd5ac0769d1e74a885450e49fde8f63ca7ab41eb479a068564aaea80b837" dependencies = [ "arrow", "arrow-buffer", @@ -2372,9 +2373,9 @@ dependencies = [ [[package]] name = "datafusion-expr" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ea229af8b73e566fc6ebb1c8dd42ab250f8855c5304bd173f4b50ead78b64304" +checksum = "86560fd7df4dda924f5474b588c56e7ae8e96dcf158c7327c9cfe9daff448ec5" dependencies = [ "arrow", "arrow-schema", @@ -2388,15 +2389,16 @@ dependencies = [ "datafusion-physical-expr-common", "indexmap 2.14.2", "itertools 0.15.0", + "serde", "serde_json", "sqlparser", ] [[package]] name = "datafusion-expr-common" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a18baa840f77760dcc9eaad281329bc57536c5c39e65e8649596f2a51216be7" +checksum = "e82e313cf655cf6bc2dce6cbd23bcb0d20236a54a1b09c436615e4bac4d8c49f" dependencies = [ "arrow", "datafusion-common", @@ -2406,9 +2408,9 @@ dependencies = [ [[package]] name = "datafusion-functions" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "763bbfeec40b55a8a961d7349c0fcfcec4f3f66cc234d6295250b95371b85e34" +checksum = "939ae2771da8a3fdc23ab4de610ae8fef9515476626b49d96d0d5c782b105771" dependencies = [ "arrow", "arrow-buffer", @@ -2438,9 +2440,9 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46c0c2b2e856883ecf6a4fa8b23e7c6041183af8a68d0583dc95d1d21e4eaee0" +checksum = "fd6c3aee9f2bc65716787d651e023743a90b4efbc57a2f679f036632e16fb110" dependencies = [ "arrow", "datafusion-common", @@ -2459,9 +2461,9 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e35e42c87dfbff06934ea0e99cfed3440ac614c3f3bfb8c33be35ca0587d4140" +checksum = "988b5d3dfb4939d81c2f9fab267063610fe52a843620f2fa6cb2a1598e43369c" dependencies = [ "arrow", "datafusion-common", @@ -2471,9 +2473,9 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08079b0604cdcc43c7d6cf006cfd5d3ed9c235ade0c1ebe91085a62f875603df" +checksum = "08149d433a97c9ea3cdf80de06480ab5c7f5453485f1e452d5b3583c81d5c0dd" dependencies = [ "arrow", "arrow-ord", @@ -2496,9 +2498,9 @@ dependencies = [ [[package]] name = "datafusion-functions-table" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fe2070ecf5638a93c81c8babc7858339ef52165a9f54da1411afd69adf7d85ce" +checksum = "2c767567cce2782d4d0bdfa55fab010aaba0bdbe533509e9e0b9a6a3ee29ed59" dependencies = [ "arrow", "async-trait", @@ -2512,9 +2514,9 @@ dependencies = [ [[package]] name = "datafusion-functions-window" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0e9a164ebe1998c5c3f9e15c9a0c4e6b599a0c8739b5471bfb089752441963f1" +checksum = "0a311c1f558d1f792f1f631c3670da12ee72a20fbf5d6ca151325d66191c03c1" dependencies = [ "arrow", "datafusion-common", @@ -2529,9 +2531,9 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "98965b183f92201c16798994820ae032649c4821e089c4c80e7d0f4485d41df5" +checksum = "fffd54b8acf6ecef284cc0677e01e1b8c71be64fec50612ba00a0ff0d9b41a52" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -2539,9 +2541,9 @@ dependencies = [ [[package]] name = "datafusion-macros" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a711bd5a06ec85d9683e142d28a98ad0679a7220669f5c1b8589e57e9b29e46" +checksum = "14d6df9763c8fb968b2417d797aef4c55ac2e4a895aaef38c72c16a287aa4a40" dependencies = [ "datafusion-doc", "quote", @@ -2550,9 +2552,9 @@ dependencies = [ [[package]] name = "datafusion-optimizer" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff3256b62ff44616ec83c7f04741d21124970833bb377ed56d9abc5ec6f707ec" +checksum = "acfb8cea6e05e2e75bae183340cf7d53cccb3db0449e2f1fbc07654d980fd19f" dependencies = [ "arrow", "chrono", @@ -2569,9 +2571,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c1e37a25b5aa1f3a1db114f0ee3605bf01639b83a831dcde35119ee61f0b11d6" +checksum = "9f002cef849e10c0597a217ec7891c2796c0f8b42d739f5d24abc24496373406" dependencies = [ "arrow", "arrow-schema", @@ -2591,9 +2593,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a4d12787d42e88b26b0d3ec006192fa23ceb832ef2a9e87439ce820a8e038cc" +checksum = "82761522c4536e7059d3b185cd99f10525c2bf48f1bd26a0a57e73e19eadc0fb" dependencies = [ "arrow", "datafusion-common", @@ -2606,9 +2608,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d13bf7f50fd902f5afc46d50759f5185b5620c40b24054f80f7c648897d3e5a6" +checksum = "f8c5c67aa252886f4d6b331c8114797f5b140464e774c4da36889269a1bf1bd8" dependencies = [ "arrow", "chrono", @@ -2623,9 +2625,9 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11fbe7e6665072957d44ed9587676e17a1f9ab3acfadfc287dac23750585f18a" +checksum = "f1a49dd3a76e4879012682c59dc3735282e60640c97483546a4113564a60bf8f" dependencies = [ "arrow", "datafusion-common", @@ -2642,9 +2644,9 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1265d58e5bce07d154e642a51ff43033b576a6ae40d50a29ac2b9f311004eb52" +checksum = "b1451d6817e5ac0dd21819b19949d1c420ab51d748f55dd9a96ecb5418a24c97" dependencies = [ "arrow", "arrow-data", @@ -2671,15 +2673,16 @@ dependencies = [ "num-traits", "parking_lot", "pin-project-lite", + "serde", "serde_json", "tokio", ] [[package]] name = "datafusion-pruning" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "74b06b333405c05015836ed36cea02cc29d1e2ee8fb414a77532d8346fb56492" +checksum = "f11f0dbe2b5a642f6106e44de8bbd6654fa7220ca14f32313d4efd5514a86968" dependencies = [ "arrow", "datafusion-common", @@ -2693,9 +2696,9 @@ dependencies = [ [[package]] name = "datafusion-session" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e7616b17b1d2fee3e9b13c651eac91b768d153265e6ba88e06e0b4af7883f69c" +checksum = "62dc5ce37a5ce1ff5362949776829aded2529a11dbaa13d24b0ffa6a059e106a" dependencies = [ "arrow-schema", "async-trait", @@ -2708,23 +2711,22 @@ dependencies = [ [[package]] name = "datafusion-spark" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f7809ecefb16fcb950c6b22ee3df99cb412dca6fae72c329d64c938ae53254f9" +checksum = "0c3d6f76463984c4ed8587c52ed4f21f142fd96fa55d8f128f673f2db4c46942" dependencies = [ "arrow", "bigdecimal", "chrono", "crc32fast", "datafusion", - "datafusion-catalog", "datafusion-common", - "datafusion-execution", "datafusion-expr", "datafusion-functions", "datafusion-functions-aggregate", "datafusion-functions-aggregate-common", "datafusion-functions-nested", + "datafusion-session", "log", "num-traits", "percent-encoding", @@ -2738,9 +2740,9 @@ dependencies = [ [[package]] name = "datafusion-sql" -version = "55.1.0" +version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b4848be66dd008d55adbe1dbc8cbd6a2502588ad0c0b8483bab0f43fd4495d0" +checksum = "816c224a0ef5469a32a562b5d4830544f714f97c5f50142109286d40a288c7db" dependencies = [ "arrow", "bigdecimal", @@ -5960,9 +5962,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.44" +version = "0.23.45" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6725596c3f2c3a0aef021139e145d4eafe314a6623e4680ca83852b2c67ab2ba" +checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634" dependencies = [ "aws-lc-rs", "once_cell", diff --git a/native/Cargo.toml b/native/Cargo.toml index a2e01fd75fa..defbf41ae27 100644 --- a/native/Cargo.toml +++ b/native/Cargo.toml @@ -44,10 +44,11 @@ arrow-select = { version = "59.2.0" } async-trait = { version = "0.1" } bytes = { version = "1.11.1" } parquet = { version = "59.2.0", default-features = false, features = ["experimental"] } -datafusion = { version = "55.1.0", default-features = false, features = ["unicode_expressions", "crypto_expressions", "nested_expressions", "parquet"] } -datafusion-datasource = { version = "55.1.0" } -datafusion-physical-expr-adapter = { version = "55.1.0" } -datafusion-spark = { version = "55.1.0", features = ["core"] } +datafusion = { version = "55.2.0", default-features = false, features = ["unicode_expressions", "crypto_expressions", "nested_expressions", "parquet"] } +datafusion-datasource = { version = "55.2.0" } +datafusion-functions = { version = "55.2.0", default-features = false, features = ["datetime_expressions"] } +datafusion-physical-expr-adapter = { version = "55.2.0" } +datafusion-spark = { version = "55.2.0", features = ["core"] } datafusion-comet-spark-expr = { path = "spark-expr" } datafusion-comet-common = { path = "common" } datafusion-comet-jni-bridge = { path = "jni-bridge" } diff --git a/native/common/src/error.rs b/native/common/src/error.rs index 1346b398061..23f92b46393 100644 --- a/native/common/src/error.rs +++ b/native/common/src/error.rs @@ -215,6 +215,9 @@ pub enum SparkError { group_index: i32, }, + #[error("[INVALID_URL] The url is invalid: {url}. If necessary set \"spark.sql.ansi.enabled\" to \"false\" to bypass this error.")] + InvalidUrl { url: String }, + #[error("[DATATYPE_CANNOT_ORDER] Cannot order by type: {data_type}.")] DatatypeCannotOrder { data_type: String }, @@ -379,6 +382,7 @@ impl SparkError { SparkError::UnexpectedPositiveValue { .. } => "UnexpectedPositiveValue", SparkError::UnexpectedNegativeValue { .. } => "UnexpectedNegativeValue", SparkError::InvalidRegexGroupIndex { .. } => "InvalidRegexGroupIndex", + SparkError::InvalidUrl { .. } => "InvalidUrl", SparkError::DatatypeCannotOrder { .. } => "DatatypeCannotOrder", SparkError::ScalarSubqueryTooManyRows => "ScalarSubqueryTooManyRows", SparkError::MergeCardinalityViolation => "MergeCardinalityViolation", @@ -613,6 +617,11 @@ impl SparkError { "groupIndex": group_index, }) } + SparkError::InvalidUrl { url } => { + serde_json::json!({ + "url": url, + }) + } SparkError::DatatypeCannotOrder { data_type } => { serde_json::json!({ "dataType": data_type, @@ -742,6 +751,7 @@ impl SparkError { // IllegalArgumentException SparkError::DatatypeCannotOrder { .. } + | SparkError::InvalidUrl { .. } | SparkError::InvalidUtf8String { .. } | SparkError::IllegalDayOfWeek { .. } | SparkError::SequenceIllegalBoundaries { .. } => { @@ -853,6 +863,9 @@ impl SparkError { // Regex errors SparkError::InvalidRegexGroupIndex { .. } => Some("INVALID_PARAMETER_VALUE"), + // URL errors + SparkError::InvalidUrl { .. } => Some("INVALID_URL"), + // Unsupported operation errors SparkError::DatatypeCannotOrder { .. } => Some("DATATYPE_CANNOT_ORDER"), diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml index 2de67226df0..b21c5bc75a8 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -111,7 +111,7 @@ jni = { version = "0.22.4", features = ["invocation"] } lazy_static = "1.4" assertables = "10" hex = "0.4.3" -datafusion-functions-nested = { version = "55.1.0" } +datafusion-functions-nested = { version = "55.2.0" } comet-test-udfs = { path = "../comet-test-udfs" } # Test-only helpers for a hand-rolled in-memory HTTP client that exercises the web-identity STS # retry path without a socket, plus reqsign's SigV4 signer for the signing round-trip test. Already diff --git a/native/core/benches/array_element_append.rs b/native/core/benches/array_element_append.rs index 2c46f9ba117..0d4f2850af5 100644 --- a/native/core/benches/array_element_append.rs +++ b/native/core/benches/array_element_append.rs @@ -22,14 +22,24 @@ //! used when processing List/Array columns in JVM shuffle. use arrow::array::builder::{ - Date32Builder, Float64Builder, Int32Builder, Int64Builder, TimestampMicrosecondBuilder, + make_builder, Date32Builder, Float64Builder, Int32Builder, Int64Builder, ListBuilder, + TimestampMicrosecondBuilder, }; use arrow::datatypes::{DataType, TimeUnit}; -use comet::execution::shuffle::spark_unsafe::list::{append_to_builder, SparkUnsafeArray}; -use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; +use comet::execution::shuffle::spark_unsafe::list::{ + append_list_element, append_to_builder, SparkUnsafeArray, +}; +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; const NUM_ELEMENTS: usize = 10000; +/// Rows in each batch of `benchmark_list_with_one_null`, as in a default Comet batch. +const LIST_ROWS: usize = 8192; + +/// Distinct arrays `benchmark_list_with_one_null` cycles through, so it does not read one array +/// from cache for every row. +const DISTINCT_ARRAYS: usize = 64; + /// Create a SparkUnsafeArray in memory with i32 elements. /// Layout: /// - 8 bytes: num_elements (i64) @@ -137,6 +147,29 @@ fn create_spark_unsafe_array_f64(num_elements: usize, with_nulls: bool) -> Vec Vec { + let header_size = 8 + num_elements.div_ceil(64) * 8; + let mut buffer = vec![0u8; header_size + (num_elements * width).div_ceil(8) * 8]; + buffer[0..8].copy_from_slice(&(num_elements as i64).to_le_bytes()); + buffer[8 + null_idx / 8] |= 1 << (null_idx % 8); + for i in 0..num_elements { + let offset = header_size + i * width; + buffer[offset..offset + width].copy_from_slice(&(i as i64).to_le_bytes()[..width]); + } + buffer + .as_chunks::<8>() + .0 + .iter() + .map(|word| u64::from_ne_bytes(*word)) + .collect() +} + fn benchmark_array_conversion(c: &mut Criterion) { let mut group = c.benchmark_group("spark_unsafe_array_to_arrow"); @@ -260,6 +293,53 @@ fn benchmark_array_conversion(c: &mut Criterion) { group.finish(); } +/// Appends a batch of list rows whose arrays each hold one null, as JVM shuffle appends a list +/// column, for array lengths on both sides of `MIN_BULK_NULLABLE_APPEND_ELEMENTS` in +/// `native/shuffle/src/spark_unsafe/list.rs`. Arrays at least that long are appended with one copy +/// of their values and a validity buffer, shorter ones element by element. To compare the two at +/// every length, save a baseline with that constant set to `usize::MAX`, which keeps every array +/// on the per-element loop, then restore it and compare against the baseline: +/// +/// ```shell +/// cargo bench --bench array_element_append -- list_with_one_null --save-baseline loop +/// cargo bench --bench array_element_append -- list_with_one_null --baseline loop +/// ``` +fn benchmark_list_with_one_null(c: &mut Criterion) { + let mut group = c.benchmark_group("list_with_one_null"); + group.throughput(Throughput::Elements(LIST_ROWS as u64)); + + for (name, element_type, width) in [("i32", DataType::Int32, 4), ("i64", DataType::Int64, 8)] { + for num_elements in [32, 64, 128] { + let buffers: Vec> = (0..DISTINCT_ARRAYS) + .map(|i| { + create_spark_unsafe_array_with_one_null(num_elements, width, i % num_elements) + }) + .collect(); + let arrays: Vec = buffers + .iter() + .map(|buffer| SparkUnsafeArray::new(buffer.as_ptr() as i64)) + .collect(); + + group.bench_with_input( + BenchmarkId::new(name, num_elements), + &arrays, + |b, arrays| { + b.iter(|| { + // The capacity Comet gives a list column's values builder. + let mut builder = ListBuilder::new(make_builder(&element_type, 100)); + for array in arrays.iter().cycle().take(LIST_ROWS) { + append_list_element(&element_type, &mut builder, array).unwrap(); + } + builder.finish() + }); + }, + ); + } + } + + group.finish(); +} + fn config() -> Criterion { Criterion::default() } @@ -267,6 +347,6 @@ fn config() -> Criterion { criterion_group! { name = benches; config = config(); - targets = benchmark_array_conversion + targets = benchmark_array_conversion, benchmark_list_with_one_null } criterion_main!(benches); diff --git a/native/core/benches/parquet_timestamp_conversion.rs b/native/core/benches/parquet_timestamp_conversion.rs index d090fc79691..6d6e6209637 100644 --- a/native/core/benches/parquet_timestamp_conversion.rs +++ b/native/core/benches/parquet_timestamp_conversion.rs @@ -119,7 +119,7 @@ fn timestamp_containers() -> [(ArrayRef, DataType); 2] { } fn benchmark(c: &mut Criterion) { - let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); let mut group = c.benchmark_group("parquet_timestamp_conversion"); for width in [8, 1024] { let (array, target) = array_sibling(width); diff --git a/native/core/src/execution/expressions/map_entries.rs b/native/core/src/execution/expressions/map_entries.rs new file mode 100644 index 00000000000..8188a80730d --- /dev/null +++ b/native/core/src/execution/expressions/map_entries.rs @@ -0,0 +1,119 @@ +// 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 std::fmt::{Display, Formatter}; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use arrow::array::{Array, ListArray, MapArray, RecordBatch}; +use arrow::datatypes::{DataType, FieldRef, Schema}; +use datafusion::common::{exec_err, Result as DataFusionResult}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::ColumnarValue; + +/// Exposes a map's entries as a list for `ExplodeExec`, sharing the input buffers. +/// Preserve the original entry fields, including nullability and metadata; the +/// SQL `map_entries` function rebuilds those fields and loses their metadata. +#[derive(Debug, Clone)] +pub struct MapEntriesExpr { + child: Arc, +} + +impl MapEntriesExpr { + pub fn new(child: Arc) -> Self { + Self { child } + } +} + +impl Display for MapEntriesExpr { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "map_entries({})", self.child) + } +} + +impl PartialEq for MapEntriesExpr { + fn eq(&self, other: &Self) -> bool { + self.child.eq(&other.child) + } +} + +impl Eq for MapEntriesExpr {} + +impl Hash for MapEntriesExpr { + fn hash(&self, state: &mut H) { + self.child.hash(state); + } +} + +impl PhysicalExpr for MapEntriesExpr { + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + Display::fmt(self, f) + } + + fn return_field(&self, input_schema: &Schema) -> DataFusionResult { + let field = self.child.return_field(input_schema)?; + let DataType::Map(entries, _) = field.data_type() else { + return exec_err!( + "MapEntriesExpr expected Map input, got {}", + field.data_type() + ); + }; + Ok(Arc::new( + field + .as_ref() + .clone() + .with_data_type(DataType::List(Arc::clone(entries))), + )) + } + + fn evaluate(&self, batch: &RecordBatch) -> DataFusionResult { + let array = self.child.evaluate(batch)?.into_array(batch.num_rows())?; + let Some(map) = array.as_any().downcast_ref::() else { + return exec_err!( + "MapEntriesExpr expected Map input, got {}", + array.data_type() + ); + }; + let DataType::Map(entries, _) = map.data_type() else { + unreachable!("MapArray downcast guarantees DataType::Map"); + }; + let list = ListArray::try_new( + Arc::clone(entries), + map.offsets().clone(), + Arc::new(map.entries().clone()), + map.nulls().cloned(), + )?; + Ok(ColumnarValue::Array(Arc::new(list))) + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.child] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> DataFusionResult> { + if children.len() != 1 { + return exec_err!( + "MapEntriesExpr expects exactly 1 child, got {}", + children.len() + ); + } + Ok(Arc::new(Self::new(Arc::clone(&children[0])))) + } +} diff --git a/native/core/src/execution/expressions/mod.rs b/native/core/src/execution/expressions/mod.rs index ad5c06a93b5..b180cf41bc0 100644 --- a/native/core/src/execution/expressions/mod.rs +++ b/native/core/src/execution/expressions/mod.rs @@ -21,6 +21,7 @@ pub mod arithmetic; pub mod bitwise; pub mod comparison; pub mod logical; +pub mod map_entries; pub mod nullcheck; pub mod partition; pub mod random; diff --git a/native/core/src/execution/expressions/temporal.rs b/native/core/src/execution/expressions/temporal.rs index 3c27c3b8263..61a9c36edd9 100644 --- a/native/core/src/execution/expressions/temporal.rs +++ b/native/core/src/execution/expressions/temporal.rs @@ -158,7 +158,12 @@ impl ExpressionBuilder for TruncTimestampBuilder { let format = planner.create_expr(expr.format.as_ref().unwrap(), input_schema)?; let timezone = expr.timezone.clone(); - Ok(Arc::new(TimestampTruncExpr::new(child, format, timezone))) + Ok(Arc::new(TimestampTruncExpr::new( + child, + format, + timezone, + expr.wrap_second_millisecond_overflow, + ))) } } diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 361d3847da7..0812b9f7087 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -24,7 +24,10 @@ //! //! `createPlan`, `setShufflePartitionPusher`, `executePlan` and `releasePlan` also depend on //! upcalls into the JVM (input iterators, the task memory manager, metrics, UDFs and scalar -//! subqueries), so their core logic keeps holding JNI references. +//! subqueries). Their core functions (`create_execution_context`, +//! `register_shuffle_partition_pusher`, `execute_plan` and `release_plan`) take the JVM objects +//! as JNI global references held by the `ExecutionContext` (`PlanJvmRefs`), and the steps that +//! need an `Env` as callbacks: the memory pool factory, the pusher factory, and `PublishMetrics`. use super::{serde, utils::SparkArrowConvert}; use crate::{ @@ -40,6 +43,7 @@ use std::collections::HashSet; use arrow::array::{Array, RecordBatch}; use arrow::datatypes::DataType as ArrowDataType; use datafusion::common::{DataFusionError, Result as DataFusionResult}; +use datafusion::config::SpillCompression; use datafusion::execution::disk_manager::DiskManagerMode; use datafusion::execution::memory_pool::MemoryPool; use datafusion::execution::runtime_env::RuntimeEnvBuilder; @@ -114,7 +118,7 @@ use tokio::sync::mpsc; use tokio::task::JoinHandle; use crate::execution::memory_pools::{ - create_memory_pool, overcommit, parse_memory_pool_config, PlanMemoryPool, + create_memory_pool, overcommit, parse_memory_pool_config, MemoryPoolConfig, PlanMemoryPool, }; use crate::execution::operators::{ScanExec, ShuffleScanExec}; use crate::execution::shuffle::{ @@ -129,9 +133,9 @@ use crate::execution::tracing::{ use crate::execution::memory_pools::logging_pool::LoggingMemoryPool; use crate::execution::spark_config::{ SparkConfig, COMET_DEBUG_ENABLED, COMET_DEBUG_MEMORY, - COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED, COMET_EXPLAIN_NATIVE_ENABLED, - COMET_MAX_TEMP_DIRECTORY_SIZE, COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, - COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, + COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED, COMET_EXEC_SPILL_COMPRESSION_CODEC, + COMET_EXPLAIN_NATIVE_ENABLED, COMET_MAX_TEMP_DIRECTORY_SIZE, + COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, }; use crate::parquet::encryption_support::{CometEncryptionFactory, ENCRYPTION_FACTORY_ID}; use crate::parquet::parquet_support::CometObjectStoreRegistry; @@ -538,8 +542,9 @@ struct ExecutionContext { batch_producer: Option, /// The pool every reservation the plan makes goes through plan_memory: Arc, - /// Native metrics - pub metrics: Arc>>, + /// The JVM `CometMetricNode` the plan's metrics are published to. `None` only without a + /// JVM (unit tests). + pub metrics: Option>>>, // The interval in milliseconds to update metrics pub metrics_update_interval: Option, // The last update time of metrics @@ -606,33 +611,16 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( class_loader_obj: JObject, ) -> jlong { try_unwrap_or_throw(&e, |env| { - // Deserialize Spark configs let bytes = env.convert_byte_array(serialized_spark_configs)?; - let spark_configs = serde::deserialize_config(bytes.as_slice())?; - let spark_config: HashMap = spark_configs.entries.into_iter().collect(); - - // Initialize the tokio runtime with spark.executor.cores as the default - // worker thread count, falling back to 1 if not set. - let executor_cores = spark_config.get_usize(SPARK_EXECUTOR_CORES, 1); - init_runtime(executor_cores); - - // Access Comet configs - let debug_native = spark_config.get_bool(COMET_DEBUG_ENABLED); - let explain_native = spark_config.get_bool(COMET_EXPLAIN_NATIVE_ENABLED); - let tracing_enabled = spark_config.get_bool(COMET_TRACING_ENABLED); - let max_temp_directory_size = - spark_config.get_u64(COMET_MAX_TEMP_DIRECTORY_SIZE, 100 * 1024 * 1024 * 1024); - let logging_memory_pool = spark_config.get_bool(COMET_DEBUG_MEMORY); - - with_trace("createPlan", tracing_enabled, || { + let settings = prepare_plan_settings(&bytes)?; + + with_trace("createPlan", settings.tracing_enabled, || { // Init JVM classes JVMClasses::init(env); let start = Instant::now(); - // Deserialize query plan - let bytes = env.convert_byte_array(serialized_query)?; - let spark_plan = serde::deserialize_op(bytes.as_slice())?; + let serialized_plan = env.convert_byte_array(serialized_query)?; let metrics = Arc::new(jni_new_global_ref!(env, metrics_node)?); @@ -645,40 +633,10 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( input_sources.push(input_source); } - // Create DataFusion memory pool let task_memory_manager = Arc::new(jni_new_global_ref!(env, comet_task_memory_manager_obj)?); let memory_pool_type = memory_pool_type.try_to_string(env)?; - let memory_pool_config = parse_memory_pool_config( - off_heap_mode != JNI_FALSE, - memory_pool_type, - memory_limit, - )?; - let memory_pool = - create_memory_pool(&memory_pool_config, task_memory_manager, task_attempt_id); - - // Register the shared base pool before wrapping it for per-plan debug logging. The - // guard removes the entry if any later plan setup step fails. - // - // Registration is not conditional on this plan's tracing setting. `tracing.enabled` is - // a session config, so an executor can run a traced and an untraced plan at once, - // while the allocation counter a trace compares against is process-wide. Registering - // only traced plans would leave the untraced plan's reservation out of the total and - // report it as allocation held outside any pool. - let rust_thread_id = get_thread_id(); - let memory_pool_registration = Some(ThreadMemoryPoolRegistration::new( - rust_thread_id, - id, - Arc::clone(&memory_pool), - )); - - let memory_pool = if logging_memory_pool { - Arc::new(LoggingMemoryPool::new(task_attempt_id as u64, memory_pool)) - } else { - memory_pool - }; - let plan_memory = Arc::new(PlanMemoryPool::new(memory_pool)); // Get local directories for storing spill files let num_local_dirs = local_dirs.len(env)?; @@ -690,46 +648,13 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( local_dirs_vec.push(local_dir); } - // We need to keep the session context alive. Some session state like temporary - // dictionaries are stored in session context. If it is dropped, the temporary - // dictionaries will be dropped as well. - let session = prepare_datafusion_session_context( - batch_size as usize, - Arc::clone(&plan_memory) as Arc, - local_dirs_vec, - max_temp_directory_size, - task_cpus as usize, - &spark_config, - &spark_plan, - )?; - - let plan_creation_time = start.elapsed(); - - let metrics_update_interval = if metrics_update_interval > 0 { - Some(Duration::from_millis(metrics_update_interval as u64)) + // Key unwrapper for encrypted files + let key_unwrapper = if !key_unwrapper_obj.is_null() { + Some(Arc::new(jni_new_global_ref!(env, key_unwrapper_obj)?)) } else { None }; - // Handle key unwrapper for encrypted files - if !key_unwrapper_obj.is_null() { - let encryption_factory = CometEncryptionFactory { - key_unwrapper: Arc::new(jni_new_global_ref!(env, key_unwrapper_obj)?), - }; - session.runtime_env().register_parquet_encryption_factory( - ENCRYPTION_FACTORY_ID, - Arc::new(encryption_factory), - ); - } - - let session = Arc::new(session); - - let tracing_event_name = if tracing_enabled { - build_tracing_event_name(&spark_plan) - } else { - String::new() - }; - // Capture the driving Spark task's TaskContext and context ClassLoader as JNI global // references when non-null. The `Arc>` releases its global ref on // drop, so cleanup is automatic when the ExecutionContext drops. @@ -745,41 +670,225 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( None }; - let exec_context = Box::new(ExecutionContext { + let inputs = PlanInputs { id, - spark_plan, + settings, + serialized_plan, partition_count: partition_count as usize, - root_op: None, - scans: vec![], - shuffle_scans: vec![], - input_sources, - stream: None, - batch_producer: None, - plan_memory, - metrics, + batch_size: batch_size as usize, + off_heap_mode: off_heap_mode != JNI_FALSE, + memory_pool_type, + memory_limit, + task_attempt_id, + task_cpus: task_cpus as usize, + local_dirs: local_dirs_vec, metrics_update_interval, - metrics_last_update_time: Instant::now(), - plan_creation_time, - session_ctx: session, - debug_native, - explain_native, - tracing_enabled, - rust_thread_id, - tracing_memory_metric_name: format!( - "thread_{rust_thread_id}_comet_memory_reserved" - ), - tracing_event_name, - task_context, - class_loader, - shuffle_partition_pusher: None, - memory_pool_registration, - }); + start, + jvm: PlanJvmRefs { + metrics: Some(metrics), + input_sources, + key_unwrapper, + task_context, + class_loader, + }, + }; + // The pool acquires memory from Spark's task memory manager. + let exec_context = create_execution_context(inputs, |config| { + create_memory_pool(config, task_memory_manager, task_attempt_id) + })?; Ok(Box::into_raw(exec_context) as i64) }) }) } +/// What `createPlan` takes from the Spark configs. +struct PlanSettings { + spark_config: HashMap, + debug_native: bool, + explain_native: bool, + tracing_enabled: bool, + max_temp_directory_size: u64, + logging_memory_pool: bool, +} + +/// Deserializes the Spark configs of a plan and initializes the Tokio runtime from them. The +/// first step of `Native.createPlan`, before its trace span opens. +fn prepare_plan_settings(serialized_spark_configs: &[u8]) -> CometResult { + let spark_configs = serde::deserialize_config(serialized_spark_configs)?; + let spark_config: HashMap = spark_configs.entries.into_iter().collect(); + + // Initialize the tokio runtime with spark.executor.cores as the default + // worker thread count, falling back to 1 if not set. + let executor_cores = spark_config.get_usize(SPARK_EXECUTOR_CORES, 1); + init_runtime(executor_cores); + + // Access Comet configs + Ok(PlanSettings { + debug_native: spark_config.get_bool(COMET_DEBUG_ENABLED), + explain_native: spark_config.get_bool(COMET_EXPLAIN_NATIVE_ENABLED), + tracing_enabled: spark_config.get_bool(COMET_TRACING_ENABLED), + max_temp_directory_size: spark_config + .get_u64(COMET_MAX_TEMP_DIRECTORY_SIZE, 100 * 1024 * 1024 * 1024), + logging_memory_pool: spark_config.get_bool(COMET_DEBUG_MEMORY), + spark_config, + }) +} + +/// The arguments of `Native.createPlan`, converted from JNI. +struct PlanInputs { + id: i64, + settings: PlanSettings, + serialized_plan: Vec, + partition_count: usize, + batch_size: usize, + off_heap_mode: bool, + memory_pool_type: String, + memory_limit: i64, + task_attempt_id: i64, + task_cpus: usize, + local_dirs: Vec, + /// In milliseconds; zero or less publishes metrics only when the plan is released. + metrics_update_interval: i64, + /// When plan creation started, the start of `ExecutionContext::plan_creation_time`. + start: Instant, + jvm: PlanJvmRefs, +} + +/// The JVM objects a plan calls back into, as JNI global references. Each is `None` or empty +/// only without a JVM (unit tests) or when Spark passes none. +struct PlanJvmRefs { + metrics: Option>>>, + input_sources: Vec>>>, + key_unwrapper: Option>>>, + task_context: Option>>>, + class_loader: Option>>>, +} + +/// Creates the execution context of a native plan: deserializes the plan, sets up its memory +/// pool and DataFusion session, and registers the pool for tracing. The plan's operators are +/// created by the first `execute_plan`. Core of `Native.createPlan`. +/// +/// `memory_pool_for` builds the pool every reservation of the plan goes through from its +/// config; `Native.createPlan` builds one that acquires memory from Spark. +fn create_execution_context( + inputs: PlanInputs, + memory_pool_for: impl FnOnce(&MemoryPoolConfig) -> Arc, +) -> CometResult> { + let PlanInputs { + id, + settings, + serialized_plan, + partition_count, + batch_size, + off_heap_mode, + memory_pool_type, + memory_limit, + task_attempt_id, + task_cpus, + local_dirs, + metrics_update_interval, + start, + jvm, + } = inputs; + + // Deserialize query plan + let spark_plan = serde::deserialize_op(serialized_plan.as_slice())?; + + // Create DataFusion memory pool + let memory_pool_config = + parse_memory_pool_config(off_heap_mode, memory_pool_type, memory_limit)?; + let memory_pool = memory_pool_for(&memory_pool_config); + + // Register the shared base pool before wrapping it for per-plan debug logging. The + // guard removes the entry if any later plan setup step fails. + // + // Registration is not conditional on this plan's tracing setting. `tracing.enabled` is + // a session config, so an executor can run a traced and an untraced plan at once, + // while the allocation counter a trace compares against is process-wide. Registering + // only traced plans would leave the untraced plan's reservation out of the total and + // report it as allocation held outside any pool. + let rust_thread_id = get_thread_id(); + let memory_pool_registration = Some(ThreadMemoryPoolRegistration::new( + rust_thread_id, + id, + Arc::clone(&memory_pool), + )); + + let memory_pool = if settings.logging_memory_pool { + Arc::new(LoggingMemoryPool::new(task_attempt_id as u64, memory_pool)) + } else { + memory_pool + }; + let plan_memory = Arc::new(PlanMemoryPool::new(memory_pool)); + + // We need to keep the session context alive. Some session state like temporary + // dictionaries are stored in session context. If it is dropped, the temporary + // dictionaries will be dropped as well. + let session = prepare_datafusion_session_context( + batch_size, + Arc::clone(&plan_memory) as Arc, + local_dirs, + settings.max_temp_directory_size, + task_cpus, + &settings.spark_config, + &spark_plan, + )?; + + let plan_creation_time = start.elapsed(); + + let metrics_update_interval = if metrics_update_interval > 0 { + Some(Duration::from_millis(metrics_update_interval as u64)) + } else { + None + }; + + // Handle key unwrapper for encrypted files + if let Some(key_unwrapper) = jvm.key_unwrapper { + let encryption_factory = CometEncryptionFactory { key_unwrapper }; + session.runtime_env().register_parquet_encryption_factory( + ENCRYPTION_FACTORY_ID, + Arc::new(encryption_factory), + ); + } + + let session = Arc::new(session); + + let tracing_event_name = if settings.tracing_enabled { + build_tracing_event_name(&spark_plan) + } else { + String::new() + }; + + Ok(Box::new(ExecutionContext { + id, + spark_plan, + partition_count, + root_op: None, + scans: vec![], + shuffle_scans: vec![], + input_sources: jvm.input_sources, + stream: None, + batch_producer: None, + plan_memory, + metrics: jvm.metrics, + metrics_update_interval, + metrics_last_update_time: Instant::now(), + plan_creation_time, + session_ctx: session, + debug_native: settings.debug_native, + explain_native: settings.explain_native, + tracing_enabled: settings.tracing_enabled, + rust_thread_id, + tracing_memory_metric_name: format!("thread_{rust_thread_id}_comet_memory_reserved"), + tracing_event_name, + task_context: jvm.task_context, + class_loader: jvm.class_loader, + shuffle_partition_pusher: None, + memory_pool_registration, + })) +} + /// Binds one task-owned shuffle callback before native execution is initialized. /// /// Keeping callback registration separate preserves the existing `createPlan` JNI ABI for @@ -799,23 +908,36 @@ pub extern "system" fn Java_org_apache_comet_Native_setShufflePartitionPusher( } let exec_context = get_execution_context(exec_context); - if exec_context.root_op.is_some() { - return Err(CometError::Internal( - "Remote shuffle callback cannot be registered after native execution starts" - .to_string(), - )); - } + register_shuffle_partition_pusher(exec_context, || { + Ok(Arc::new(JavaShufflePartitionPusher::try_new( + env, &callback, + )?)) + }) + }) +} - if exec_context.shuffle_partition_pusher.is_some() { - return Err(CometError::Internal( - "Remote shuffle callback has already been registered for this task".to_string(), - )); - } +/// Registers the plan's remote shuffle callback, built by `make_pusher` only once the plan +/// is known to accept one: before execution starts, and once. Core of +/// `Native.setShufflePartitionPusher`, whose pusher calls back into the JVM. +fn register_shuffle_partition_pusher( + exec_context: &mut ExecutionContext, + make_pusher: impl FnOnce() -> CometResult>, +) -> CometResult<()> { + if exec_context.root_op.is_some() { + return Err(CometError::Internal( + "Remote shuffle callback cannot be registered after native execution starts" + .to_string(), + )); + } - let pusher = JavaShufflePartitionPusher::try_new(env, &callback)?; - exec_context.shuffle_partition_pusher = Some(Arc::new(pusher)); - Ok(()) - }) + if exec_context.shuffle_partition_pusher.is_some() { + return Err(CometError::Internal( + "Remote shuffle callback has already been registered for this task".to_string(), + )); + } + + exec_context.shuffle_partition_pusher = Some(make_pusher()?); + Ok(()) } /// Skipping is opt-in (`spark.comet.exec.aggregate.skipPartial.enabled`): once DataFusion's probe @@ -902,6 +1024,24 @@ fn prepare_datafusion_session_context( session_config.set_str("datafusion.execution.parquet.reorder_filters", "true"); } + // Translate Comet's spill codec into DataFusion's spill compression, which DataFusion + // otherwise leaves uncompressed. Set before the pass-through below for the same reason, + // so an explicit `spark.comet.datafusion.execution.spill_compression` still wins. + let spill_compression = match spark_config + .get(COMET_EXEC_SPILL_COMPRESSION_CODEC) + .map(String::as_str) + { + None | Some("lz4") => SpillCompression::Lz4Frame, + Some("zstd") => SpillCompression::Zstd, + Some("none") => SpillCompression::Uncompressed, + Some(codec) => { + return Err(CometError::Config(format!( + "Unsupported spill compression codec: {codec}" + ))) + } + }; + session_config = session_config.with_spill_compression(spill_compression); + // Pass through DataFusion configs from Spark. // e.g: spark-shell --conf spark.comet.datafusion.sql_parser.parse_float_as_decimal=true // becomes datafusion.sql_parser.parse_float_as_decimal=true @@ -1259,93 +1399,8 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( }; let result = with_trace(tracing_label, tracing_enabled, || { - let exec_context_id = exec_context.id; - - // Initialize the execution stream. - // Because we don't know if input arrays are dictionary-encoded when we create - // query plan, we need to defer stream initialization to first time execution. - if exec_context.root_op.is_none() { - let start = Instant::now(); - let planner = - PhysicalPlanner::new(Arc::clone(&exec_context.session_ctx), partition) - .with_exec_id(exec_context_id) - .with_sql_text_pool(&exec_context.spark_plan) - .with_task_context(exec_context.task_context.clone()) - .with_class_loader(exec_context.class_loader.clone()) - .with_shuffle_partition_pusher( - exec_context.shuffle_partition_pusher.clone(), - ); - let (scans, shuffle_scans, root_op) = planner.create_plan( - &exec_context.spark_plan, - &mut exec_context.input_sources.clone(), - exec_context.partition_count, - )?; - let physical_plan_time = start.elapsed(); - - exec_context.plan_creation_time += physical_plan_time; - exec_context.scans = scans; - exec_context.shuffle_scans = shuffle_scans; - - if exec_context.explain_native { - let formatted_plan_str = - DisplayableExecutionPlan::new(root_op.native_plan.as_ref()).indent(true); - info!("Comet native query plan:\n{formatted_plan_str:}"); - } - - let task_ctx = exec_context.session_ctx.task_ctx(); - // Each Comet native execution corresponds to a single Spark partition, - // so we should always execute partition 0. - let stream = root_op.native_plan.execute(0, task_ctx)?; - - if exec_context.scans.is_empty() && exec_context.shuffle_scans.is_empty() { - // No JVM data sources — spawn onto tokio so the executor - // thread parks in blocking_recv instead of busy-polling. - exec_context.batch_producer = - Some(BatchProducer::spawn(&get_runtime(), stream)); - } else { - exec_context.stream = Some(stream); - } - exec_context.root_op = Some(root_op); - } else { - // Pull input batches - pull_input_batches(exec_context)?; - } - - if let Some(producer) = &mut exec_context.batch_producer { - match producer.next_batch()? { - Some(batch) => { - // Publish on the configured interval, as the ScanExec path below does, - // since each publish walks the whole metric tree and calls into the JVM. - // `releasePlan` publishes the final values. - update_metrics_on_interval(env, exec_context)?; - return prepare_output( - env, - array_addrs, - schema_addrs, - batch, - exec_context.debug_native, - ); - } - None => { - log_plan_metrics(exec_context, stage_id, partition); - return Ok(-1); - } - } - } - - // ScanExec path: JVM-fed scans return `Pending` until `pull_input_batches` refills - // them and wakes the stream. A poll that is still pending, with nothing having woken - // the stream by the end of the pull, waits on native I/O, and `next_batch` parks - // until it completes. - let mut stream = exec_context.stream.take().unwrap(); - let next = get_runtime().block_on(next_batch(&mut stream, || { - pull_input_batches(exec_context)?; - update_metrics_on_interval(env, exec_context) - })); - exec_context.stream = Some(stream); - let next = next?; - update_metrics_on_interval(env, exec_context)?; - match next { + let mut publish_metrics = |ctx: &ExecutionContext| update_metrics(env, ctx); + match execute_plan(exec_context, stage_id, partition, &mut publish_metrics)? { Some(batch) => prepare_output( env, array_addrs, @@ -1353,10 +1408,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( batch, exec_context.debug_native, ), - None => { - log_plan_metrics(exec_context, stage_id, partition); - Ok(-1) - } + None => Ok(-1), } }); @@ -1381,6 +1433,103 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( }) } +/// Publishes a plan's metrics to the JVM. The plan entry points pass one that calls into the JVM +/// through their `Env`, so the core functions that drive a plan take no JNI types. +type PublishMetrics<'a> = dyn FnMut(&ExecutionContext) -> CometResult<()> + 'a; + +/// Returns the plan's next output batch, or `None` once the plan has produced all of its output. +/// The first call creates the plan's operators and starts its stream; a plan with no JVM input +/// then runs on a Tokio task. Metrics are published through `publish_metrics` on the configured +/// interval. Core of `Native.executePlan`; the JVM-fed scans call back into the JVM for their +/// input. +fn execute_plan( + exec_context: &mut ExecutionContext, + stage_id: i32, + partition: i32, + publish_metrics: &mut PublishMetrics, +) -> CometResult> { + let exec_context_id = exec_context.id; + + // Initialize the execution stream. + // Because we don't know if input arrays are dictionary-encoded when we create + // query plan, we need to defer stream initialization to first time execution. + if exec_context.root_op.is_none() { + let start = Instant::now(); + let planner = PhysicalPlanner::new(Arc::clone(&exec_context.session_ctx), partition) + .with_exec_id(exec_context_id) + .with_sql_text_pool(&exec_context.spark_plan) + .with_task_context(exec_context.task_context.clone()) + .with_class_loader(exec_context.class_loader.clone()) + .with_shuffle_partition_pusher(exec_context.shuffle_partition_pusher.clone()); + let (scans, shuffle_scans, root_op) = planner.create_plan( + &exec_context.spark_plan, + &mut exec_context.input_sources.clone(), + exec_context.partition_count, + )?; + let physical_plan_time = start.elapsed(); + + exec_context.plan_creation_time += physical_plan_time; + exec_context.scans = scans; + exec_context.shuffle_scans = shuffle_scans; + + if exec_context.explain_native { + let formatted_plan_str = + DisplayableExecutionPlan::new(root_op.native_plan.as_ref()).indent(true); + info!("Comet native query plan:\n{formatted_plan_str:}"); + } + + let task_ctx = exec_context.session_ctx.task_ctx(); + // Each Comet native execution corresponds to a single Spark partition, + // so we should always execute partition 0. + let stream = root_op.native_plan.execute(0, task_ctx)?; + + if exec_context.scans.is_empty() && exec_context.shuffle_scans.is_empty() { + // No JVM data sources — spawn onto tokio so the executor + // thread parks in blocking_recv instead of busy-polling. + exec_context.batch_producer = Some(BatchProducer::spawn(&get_runtime(), stream)); + } else { + exec_context.stream = Some(stream); + } + exec_context.root_op = Some(root_op); + } else { + // Pull input batches + pull_input_batches(exec_context)?; + } + + if let Some(producer) = &mut exec_context.batch_producer { + return match producer.next_batch()? { + Some(batch) => { + // Publish on the configured interval, as the ScanExec path below does, + // since each publish walks the whole metric tree and calls into the JVM. + // `releasePlan` publishes the final values. + update_metrics_on_interval(exec_context, publish_metrics)?; + Ok(Some(batch)) + } + None => { + log_plan_metrics(exec_context, stage_id, partition); + Ok(None) + } + }; + } + + // ScanExec path: JVM-fed scans return `Pending` until `pull_input_batches` refills + // them and wakes the stream. A poll that is still pending, with nothing having woken + // the stream by the end of the pull, waits on native I/O, and `next_batch` parks + // until it completes. + let mut stream = exec_context.stream.take().unwrap(); + let next = get_runtime().block_on(next_batch(&mut stream, || { + pull_input_batches(exec_context)?; + update_metrics_on_interval(exec_context, publish_metrics) + })); + exec_context.stream = Some(stream); + let next = next?; + update_metrics_on_interval(exec_context, publish_metrics)?; + if next.is_none() { + log_plan_metrics(exec_context, stage_id, partition); + } + Ok(next) +} + #[no_mangle] /// Drop the native query plan object and context object. pub extern "system" fn Java_org_apache_comet_Native_releasePlan( @@ -1400,48 +1549,58 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan( // Reclaim ownership of the context up front so that it is always freed, even if a step // below fails. Dropping it releases the memory pool and every JNI global ref the context // holds. - let mut execution_context: Box = + let execution_context: Box = Box::from_raw(exec_context as *mut ExecutionContext); - // A plan with no JVM input runs on a Tokio task. Stop it and drop its stream first, so the - // stream's memory is returned before the Spark task can end, and the metrics flushed below - // include what the stream records as it is dropped. - let producer_stopped = execution_context - .batch_producer - .take() - .map_or(Ok(()), BatchProducer::stop); - - // Unregister this context's pool and, when tracing, emit the remaining total for the - // thread. Every context registers, but only a traced one writes counters, so the - // reservations are read only then. - drop(execution_context.memory_pool_registration.take()); - if execution_context.tracing_enabled { - log_memory_usage( - &execution_context.tracing_memory_metric_name, - total_reserved_for_thread(execution_context.rust_thread_id) as u64, - ); - } + release_plan(execution_context, &mut |ctx| update_metrics(env, ctx)) + }) +} - let metrics_flushed = update_metrics(env, &mut execution_context); - - // Once this returns the Spark task can end, and Spark hands whatever the task still holds - // to other tasks. Dropping the plan aborts the tasks its operators spawned, and they give - // back what they hold the next time they yield, so wait for them. - let id = execution_context.id; - let plan_memory = Arc::clone(&execution_context.plan_memory); - drop(execution_context); - let held = plan_memory.wait_until_released(Instant::now() + PLAN_MEMORY_RELEASE_TIMEOUT); - if held > 0 { - warn!( - "Native plan {id} still holds {held} bytes {PLAN_MEMORY_RELEASE_TIMEOUT:?} after \ - it was released. They go back to Spark only when whatever holds them drops them, \ - which can be after the Spark task has ended and its memory has been given to \ - another task." - ); - } +/// Releases a plan: stops it, publishes its final metrics through `publish_metrics`, drops it, +/// and waits briefly for its memory reservations to be returned. The context is freed even if a +/// step fails. Core of `Native.releasePlan`. +fn release_plan( + mut execution_context: Box, + publish_metrics: &mut PublishMetrics, +) -> CometResult<()> { + // A plan with no JVM input runs on a Tokio task. Stop it and drop its stream first, so the + // stream's memory is returned before the Spark task can end, and the metrics flushed below + // include what the stream records as it is dropped. + let producer_stopped = execution_context + .batch_producer + .take() + .map_or(Ok(()), BatchProducer::stop); + + // Unregister this context's pool and, when tracing, emit the remaining total for the + // thread. Every context registers, but only a traced one writes counters, so the + // reservations are read only then. + drop(execution_context.memory_pool_registration.take()); + if execution_context.tracing_enabled { + log_memory_usage( + &execution_context.tracing_memory_metric_name, + total_reserved_for_thread(execution_context.rust_thread_id) as u64, + ); + } - producer_stopped.and(metrics_flushed) - }) + let metrics_flushed = publish_metrics(&execution_context); + + // Once this returns the Spark task can end, and Spark hands whatever the task still holds + // to other tasks. Dropping the plan aborts the tasks its operators spawned, and they give + // back what they hold the next time they yield, so wait for them. + let id = execution_context.id; + let plan_memory = Arc::clone(&execution_context.plan_memory); + drop(execution_context); + let held = plan_memory.wait_until_released(Instant::now() + PLAN_MEMORY_RELEASE_TIMEOUT); + if held > 0 { + warn!( + "Native plan {id} still holds {held} bytes {PLAN_MEMORY_RELEASE_TIMEOUT:?} after \ + it was released. They go back to Spark only when whatever holds them drops them, \ + which can be after the Spark task has ended and its memory has been given to \ + another task." + ); + } + + producer_stopped.and(metrics_flushed) } /// How long `releasePlan` waits for a released plan's reservations to be returned. The tasks @@ -1450,11 +1609,11 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan( /// Spark task. const PLAN_MEMORY_RELEASE_TIMEOUT: Duration = Duration::from_secs(1); -/// Runs `update_metrics` once the configured interval has passed and, with tracing on, samples +/// Runs `publish_metrics` once the configured interval has passed and, with tracing on, samples /// this thread's pool reservation at the same cadence. fn update_metrics_on_interval( - env: &mut Env, exec_context: &mut ExecutionContext, + publish_metrics: &mut PublishMetrics, ) -> CometResult<()> { let Some(interval) = exec_context.metrics_update_interval else { return Ok(()); @@ -1463,7 +1622,7 @@ fn update_metrics_on_interval( if now - exec_context.metrics_last_update_time < interval { return Ok(()); } - update_metrics(env, exec_context)?; + publish_metrics(exec_context)?; exec_context.metrics_last_update_time = now; if exec_context.tracing_enabled { log_memory_usage( @@ -1474,12 +1633,13 @@ fn update_metrics_on_interval( Ok(()) } -fn update_metrics(env: &mut Env, exec_context: &mut ExecutionContext) -> CometResult<()> { - if let Some(native_query) = &exec_context.root_op { - let metrics = exec_context.metrics.as_obj(); - update_comet_metric(env, metrics, native_query) - } else { - Ok(()) +/// Publishes the plan's metrics to its JVM `CometMetricNode`, once the plan has been created. +fn update_metrics(env: &mut Env, exec_context: &ExecutionContext) -> CometResult<()> { + match (&exec_context.root_op, &exec_context.metrics) { + (Some(native_query), Some(metrics)) => { + update_comet_metric(env, metrics.as_obj(), native_query) + } + _ => Ok(()), } } @@ -3381,4 +3541,244 @@ mod tests { "{error}" ); } + + /// A range of `num_elements` longs from 0, which needs no JVM input and reserves no memory. + fn range_operator(num_elements: i64) -> Operator { + use datafusion_comet_proto::spark_operator::RangeScan; + + Operator { + plan_id: 1, + op_struct: Some(OpStruct::RangeScan(RangeScan { + start: 0, + step: 1, + num_elements, + num_slices: 1, + })), + ..Default::default() + } + } + + fn range_plan(num_elements: i64) -> Vec { + use prost::Message; + range_operator(num_elements).encode_to_vec() + } + + /// A range of `num_elements` longs sorted in descending order. The sort buffers its whole + /// input, so the plan holds a memory reservation while it emits its output. + fn sorted_range_plan(num_elements: i64) -> Vec { + use datafusion_comet_proto::spark_expression::{ + self, expr::ExprStruct, NullOrdering, SortDirection, + }; + use datafusion_comet_proto::spark_operator::Sort; + use prost::Message; + + let column = Expr { + expr_struct: Some(ExprStruct::Bound(spark_expression::BoundReference { + index: 0, + // A Spark `long`. + datatype: Some(spark_expression::DataType { + type_id: 4, + type_info: None, + }), + })), + ..Default::default() + }; + let descending = Expr { + expr_struct: Some(ExprStruct::SortOrder(Box::new( + spark_expression::SortOrder { + child: Some(Box::new(column)), + direction: SortDirection::Descending as i32, + null_ordering: NullOrdering::NullsLast as i32, + }, + ))), + ..Default::default() + }; + Operator { + plan_id: 2, + children: vec![range_operator(num_elements)], + op_struct: Some(OpStruct::Sort(Sort { + sort_orders: vec![descending], + ..Default::default() + })), + ..Default::default() + } + .encode_to_vec() + } + + fn long_values(batch: &RecordBatch) -> Vec { + arrow::array::AsArray::as_primitive::(batch.column(0)) + .values() + .to_vec() + } + + /// Creates a plan without a JVM: no metrics node or JVM input, and a memory pool backed by a + /// fake Spark. + fn create_test_plan(serialized_plan: Vec) -> CometResult> { + use crate::execution::memory_pools::create_memory_pool_with_fake_spark; + use datafusion_comet_proto::spark_config::ConfigMap; + use prost::Message; + + let configs = ConfigMap { + entries: HashMap::from([(SPARK_EXECUTOR_CORES.to_string(), "2".to_string())]), + }; + let inputs = PlanInputs { + id: -6201, + settings: prepare_plan_settings(&configs.encode_to_vec())?, + serialized_plan, + partition_count: 1, + batch_size: 4, + off_heap_mode: true, + memory_pool_type: "greedy_unified".to_string(), + memory_limit: 1 << 30, + task_attempt_id: -6201, + task_cpus: 1, + local_dirs: vec![std::env::temp_dir().to_str().unwrap().to_string()], + metrics_update_interval: 0, + start: Instant::now(), + jvm: PlanJvmRefs { + metrics: None, + input_sources: vec![], + key_unwrapper: None, + task_context: None, + class_loader: None, + }, + }; + create_execution_context(inputs, |config| { + create_memory_pool_with_fake_spark(config, -6201, 1 << 30) + }) + } + + #[test] + fn execute_plan_drains_a_plan() { + let _guard = serial(); + let mut exec_context = create_test_plan(range_plan(10)).unwrap(); + assert!( + exec_context.root_op.is_none(), + "operators are created on first execution" + ); + + let mut values = Vec::new(); + while let Some(batch) = execute_plan(&mut exec_context, 0, 0, &mut |_| Ok(())).unwrap() { + assert!( + batch.num_rows() <= 4, + "batches follow the configured batch size" + ); + values.extend(long_values(&batch)); + } + assert_eq!(values, (0..10).collect::>()); + assert!( + exec_context.batch_producer.is_some(), + "a plan without JVM input runs on a Tokio task" + ); + release_plan(exec_context, &mut |_| Ok(())).unwrap(); + } + + /// The sort holds a reservation while it emits its 10,000 sorted rows, 4 at a time. + const SORTED_ROWS: i64 = 10_000; + + #[test] + fn release_plan_returns_the_plans_memory() { + let _guard = serial(); + let mut exec_context = create_test_plan(sorted_range_plan(SORTED_ROWS)).unwrap(); + + // Without an update interval, metrics are published only when the plan is released. + let mut publishes = 0; + let batch = execute_plan(&mut exec_context, 0, 0, &mut |_| { + publishes += 1; + Ok(()) + }) + .unwrap() + .expect("the sort emits a first batch"); + assert_eq!(long_values(&batch), [9999, 9998, 9997, 9996]); + assert_eq!(publishes, 0); + + // Released mid-stream, as when Spark stops reading early, with the sort still holding + // its buffered input. + let plan_memory = Arc::clone(&exec_context.plan_memory); + assert!(plan_memory.reserved() > 0, "the sort holds a reservation"); + release_plan(exec_context, &mut |ctx| { + assert!( + ctx.root_op.is_some(), + "the final metrics are of the executed plan" + ); + publishes += 1; + Ok(()) + }) + .unwrap(); + assert_eq!(publishes, 1); + assert_eq!(plan_memory.reserved(), 0); + } + + #[test] + fn release_plan_reports_a_failed_metrics_publish() { + let _guard = serial(); + let mut exec_context = create_test_plan(sorted_range_plan(SORTED_ROWS)).unwrap(); + execute_plan(&mut exec_context, 0, 0, &mut |_| Ok(())).unwrap(); + let plan_memory = Arc::clone(&exec_context.plan_memory); + assert!(plan_memory.reserved() > 0, "the sort holds a reservation"); + + let error = release_plan(exec_context, &mut |_| { + Err(CometError::Internal("metrics node is gone".to_string())) + }) + .unwrap_err(); + assert!( + error.to_string().contains("metrics node is gone"), + "{error}" + ); + // The plan is still released. + assert_eq!(plan_memory.reserved(), 0); + } + + #[test] + fn create_execution_context_rejects_bad_input() { + let _guard = serial(); + let error = create_test_plan(vec![0xff, 0xff, 0xff]) + .err() + .expect("a malformed plan is rejected"); + assert!(error.to_string().contains("deserialize"), "{error}"); + } + + struct NoopPusher; + + impl ShufflePartitionPusher for NoopPusher { + fn push_partition_data( + &self, + _partition_id: i32, + _data: &[u8], + ) -> datafusion::common::Result<()> { + Ok(()) + } + } + + #[test] + fn register_shuffle_partition_pusher_once_before_execution() { + let _guard = serial(); + let mut exec_context = create_test_plan(range_plan(3)).unwrap(); + register_shuffle_partition_pusher(&mut exec_context, || Ok(Arc::new(NoopPusher))).unwrap(); + assert!(exec_context.shuffle_partition_pusher.is_some()); + + // The pusher is built only once the plan is known to accept it. + let error = register_shuffle_partition_pusher(&mut exec_context, || { + panic!("a rejected pusher must not be built") + }) + .unwrap_err(); + assert!( + error.to_string().contains("already been registered"), + "{error}" + ); + + let mut exec_context = create_test_plan(range_plan(3)).unwrap(); + execute_plan(&mut exec_context, 0, 0, &mut |_| Ok(())).unwrap(); + let error = register_shuffle_partition_pusher(&mut exec_context, || { + panic!("a rejected pusher must not be built") + }) + .unwrap_err(); + assert!( + error + .to_string() + .contains("cannot be registered after native execution starts"), + "{error}" + ); + release_plan(exec_context, &mut |_| Ok(())).unwrap(); + } } diff --git a/native/core/src/execution/memory_pools/mod.rs b/native/core/src/execution/memory_pools/mod.rs index 1fc4db2021c..e2fd2e58a73 100644 --- a/native/core/src/execution/memory_pools/mod.rs +++ b/native/core/src/execution/memory_pools/mod.rs @@ -22,6 +22,8 @@ mod plan_pool; mod spark_memory; mod spill_replay; mod task_shared; +#[cfg(test)] +pub(crate) mod testing; mod unified_pool; use datafusion::execution::memory_pool::{MemoryPool, TrackConsumersPool, UnboundedMemoryPool}; diff --git a/native/core/src/execution/memory_pools/testing.rs b/native/core/src/execution/memory_pools/testing.rs new file mode 100644 index 00000000000..ce1603a26f5 --- /dev/null +++ b/native/core/src/execution/memory_pools/testing.rs @@ -0,0 +1,81 @@ +// 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. + +//! Memory pools for tests of operators that reserve memory. + +use std::fmt; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use datafusion::common::Result; +use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool, MemoryReservation}; + +/// A [`GreedyMemoryPool`] that also records the most it ever had reserved, so a test can check +/// what an operator reserved while it ran, not only that it gave everything back. +#[derive(Debug)] +pub(crate) struct PeakMemoryPool { + inner: GreedyMemoryPool, + peak: AtomicUsize, +} + +impl PeakMemoryPool { + pub(crate) fn new(limit: usize) -> Self { + Self { + inner: GreedyMemoryPool::new(limit), + peak: AtomicUsize::new(0), + } + } + + pub(crate) fn peak(&self) -> usize { + self.peak.load(Ordering::Relaxed) + } + + fn record(&self) { + self.peak + .fetch_max(self.inner.reserved(), Ordering::Relaxed); + } +} + +impl fmt::Display for PeakMemoryPool { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "PeakMemoryPool({}, peak {})", self.inner, self.peak()) + } +} + +impl MemoryPool for PeakMemoryPool { + fn name(&self) -> &str { + "PeakMemoryPool" + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.inner.grow(reservation, additional); + self.record(); + } + + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.inner.shrink(reservation, shrink); + } + + fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> { + self.inner.try_grow(reservation, additional)?; + self.record(); + Ok(()) + } + + fn reserved(&self) -> usize { + self.inner.reserved() + } +} diff --git a/native/core/src/execution/operators/dynamic_filter/parquet_reader/schema_adapter/tests/resolution.rs b/native/core/src/execution/operators/dynamic_filter/parquet_reader/schema_adapter/tests/resolution.rs index 21b7bcc6315..706f75f7f2e 100644 --- a/native/core/src/execution/operators/dynamic_filter/parquet_reader/schema_adapter/tests/resolution.rs +++ b/native/core/src/execution/operators/dynamic_filter/parquet_reader/schema_adapter/tests/resolution.rs @@ -27,7 +27,7 @@ fn spark_factory( use_field_id: bool, defaults: Option>, ) -> Arc { - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.use_field_id = use_field_id; Arc::new(SparkPhysicalExprAdapterFactory::new(options, defaults)) } diff --git a/native/core/src/execution/operators/iceberg_common.rs b/native/core/src/execution/operators/iceberg_common.rs index f4ea3f7983b..a17eb076e0d 100644 --- a/native/core/src/execution/operators/iceberg_common.rs +++ b/native/core/src/execution/operators/iceberg_common.rs @@ -48,16 +48,12 @@ const ICEBERG_PROVIDER_CLASS_PROPERTY: &str = "s3.comet.credential.provider.clas /// iceberg-storage-opendal's own settings, such as `opendal.io-timeout-ms`. const STORAGE_PROPERTY_PREFIXES: &[&str] = &["s3.", "gcs.", "adls.", "client.", "opendal."]; -/// Pick an OpenDAL storage backend from a URI's scheme. `file` (or no scheme) falls through to -/// the local file system. `memory` is used by the write path to assemble manifest bytes that -/// stay entirely in-process. For S3, the Comet credential bridge is wired in when a provider -/// class is configured; `access_mode` is forwarded to the JVM SPI so the read and write paths can -/// be granted different (e.g. read-only vs read-write) credentials. +/// Pick an OpenDAL storage backend for a URI whose scheme `builtin_storage_schemes` lists for +/// `access_mode`, or that is an opted-in S3-compliant alias written in lowercase; anything else is +/// rejected before the match, so the list alone decides what each mode admits. For S3, the Comet +/// credential bridge is wired in when a provider class is configured and `access_mode` is +/// forwarded to the JVM SPI. /// -/// The JVM planner mirrors the schemes admitted here by hand, so that it can decline a scan -/// cleanly instead of failing at execution. Changing the arms below means updating -/// `CometScanRule.icebergReadableSchemes` (reads) and -/// `CometIcebergNativeWrite.SupportedStorageSchemes` (writes). /// The bool is false when a read fell back to opendal's default chain because the configured S3 /// access provider failed to initialise; such a `FileIO` must not be cached, so the next task /// retries. @@ -67,24 +63,22 @@ pub(crate) fn storage_factory_for( catalog_name: &str, access_mode: AccessMode, ) -> Result<(Arc, bool), DataFusionError> { + // Verbatim match: OpenDAL strips the scheme prefix from every path case-sensitively at open + // time, so admitting `S3://` here would only defer the failure. Aliases are held to the same + // rule (`is_iceberg_alias_scheme`). The JVM gates match verbatim. let scheme = scheme_of(path); + if !builtin_storage_schemes(access_mode).contains(&scheme) + && !is_s3_family_scheme(scheme, catalog_properties) + { + return Err(DataFusionError::Execution(format!( + "Unsupported storage scheme: {scheme}" + ))); + } match scheme { "file" => Ok((Arc::new(OpenDalStorageFactory::Fs), true)), "memory" => Ok((Arc::new(OpenDalStorageFactory::Memory), true)), "gs" => Ok((Arc::new(OpenDalStorageFactory::Gcs), true)), - // Reads keep the OSS backend they have always had (CometScanRule admits `oss` scan - // locations through HadoopFileIO). Writes fail closed: Comet does not forward `oss.*` - // properties into the FileIO and no test covers the write path, so OSS-specific - // endpoint/credential configuration could silently be dropped. The JVM write gate - // already declines `oss` locations; this is the native-side backstop. - "oss" => match access_mode { - AccessMode::Read => Ok((Arc::new(OpenDalStorageFactory::Oss), true)), - AccessMode::Write => Err(DataFusionError::Execution( - "OSS is not supported for native Iceberg writes (oss.* properties are not \ - forwarded to the native FileIO)" - .to_string(), - )), - }, + "oss" => Ok((Arc::new(OpenDalStorageFactory::Oss), true)), // s3, s3a, and any opted-in s3-compliant alias (e.g. blob) route to the S3 backend. Listed // last so the built-in backends above stay authoritative even if one of their schemes is // also named in `fs.comet.s3Compliant.schemes`. An alias additionally gets a wrapper that @@ -93,7 +87,7 @@ pub(crate) fn storage_factory_for( // location-scoped provider gets a storage that serves each file with the credential of // its policy location -- see iceberg_location_scoped. s if is_s3_family_scheme(s, catalog_properties) => { - let alias = is_s3_compliant_alias_scheme(s, catalog_properties); + let alias = is_iceberg_alias_scheme(s, catalog_properties); let (access, cacheable) = build_s3_access(path, catalog_properties, catalog_name, access_mode)?; let factory: Arc = match access { @@ -111,6 +105,7 @@ pub(crate) fn storage_factory_for( }; Ok((factory, cacheable)) } + // Only a listed scheme without an arm reaches here; the tests below fail on that. _ => Err(DataFusionError::Execution(format!( "Unsupported storage scheme: {scheme}" ))), @@ -303,6 +298,17 @@ fn cached_file_io( Ok(file_io) } +/// The single point of change for the built-in schemes `storage_factory_for` admits per access +/// mode; the JVM read and write gates load these lists over JNI. `memory` is write-only: an OpenDAL +/// memory backend is a fresh empty in-process store the write path assembles manifests in, so a +/// read finds nothing. `oss` is read-only: no `oss.*` property is forwarded and no test covers it. +pub(crate) fn builtin_storage_schemes(access_mode: AccessMode) -> &'static [&'static str] { + match access_mode { + AccessMode::Read => &["file", "gs", "oss", "s3", "s3a"], + AccessMode::Write => &["file", "memory", "gs", "s3", "s3a"], + } +} + pub(crate) fn load_file_io( catalog_properties: &HashMap, reference_path: &str, @@ -574,7 +580,7 @@ fn env_region_present() -> bool { /// The JVM write gate (`CometIcebergNativeWrite.storageScheme`) mirrors this rule exactly, case /// included. Change both together, and keep the cases in /// `scheme_of_extracts_scheme_from_all_uri_forms` in step with its `storageScheme` test. -fn scheme_of(path: &str) -> &str { +pub(crate) fn scheme_of(path: &str) -> &str { match path.split_once(':') { Some((scheme, _)) if !scheme.is_empty() && !scheme.contains('/') => scheme, _ => "file", @@ -587,7 +593,19 @@ fn scheme_of(path: &str) -> &str { /// iceberg-rust's storage factory has no `s3n` backend and the Scala Iceberg scheme gate /// (`isIcebergReadableScheme`) rejects `s3n` before a path ever reaches this operator. fn is_s3_family_scheme(scheme: &str, catalog_properties: &HashMap) -> bool { - matches!(scheme, "s3" | "s3a") || is_s3_compliant_alias_scheme(scheme, catalog_properties) + matches!(scheme, "s3" | "s3a") || is_iceberg_alias_scheme(scheme, catalog_properties) +} + +/// True if `scheme` is an opted-in S3-compliant alias written exactly as the lowercase form of its +/// list entry. The shared `is_s3_compliant_alias_scheme` is case-insensitive, which is safe for +/// the Parquet path because it rewrites an alias URL to `s3://` before anything opens it. The +/// Iceberg path opens the recorded location as written, and OpenDAL's S3 backend checks it against +/// a `scheme://bucket/` prefix whose scheme comes from `Url::parse` and so is lowercase, so +/// `BLOB://bucket/key` would pass a case-insensitive gate here and fail at open time. The JVM +/// Iceberg gate matches the same way. +fn is_iceberg_alias_scheme(scheme: &str, catalog_properties: &HashMap) -> bool { + !scheme.bytes().any(|b| b.is_ascii_uppercase()) + && is_s3_compliant_alias_scheme(scheme, catalog_properties) } #[cfg(test)] @@ -823,7 +841,25 @@ mod tests { // oss.* property forwarding exists and is tested. assert!(factory_result("oss://bucket/db/table", AccessMode::Read).is_ok()); let err = factory_result("oss://bucket/db/table", AccessMode::Write).unwrap_err(); - assert!(err.contains("OSS"), "unexpected error: {err}"); + assert!( + err.contains("Unsupported storage scheme: oss"), + "unexpected error: {err}" + ); + } + + #[test] + fn memory_scheme_is_writable_but_not_readable() { + // The write path assembles manifests in a fresh in-process memory store. A read against + // a new memory store can never find data, so the factory must decline it up front. + assert!(factory_result("memory:manifest.avro", AccessMode::Write).is_ok()); + assert!(factory_result("memory:///manifest.avro", AccessMode::Write).is_ok()); + for path in ["memory:manifest.avro", "memory:///key.parquet"] { + let err = factory_result(path, AccessMode::Read).unwrap_err(); + assert!( + err.contains("Unsupported storage scheme: memory"), + "unexpected error for {path}: {err}" + ); + } } #[test] @@ -831,13 +867,85 @@ mod tests { for mode in [AccessMode::Read, AccessMode::Write] { assert!(factory_result("file:///tmp/x", mode).is_ok()); assert!(factory_result("/tmp/no-scheme", mode).is_ok()); - assert!(factory_result("memory:manifest.avro", mode).is_ok()); // No credential provider configured: the default chain applies in both modes. assert!(factory_result("s3://bucket/db/table", mode).is_ok()); assert!(factory_result("gs://bucket/db/table", mode).is_ok()); } } + #[test] + fn scheme_list_pre_check_is_load_bearing() { + // The oss and memory arms build a backend for either mode; only their absence from the + // list for the other mode rejects them. Both facts are asserted so that adding an + // access-mode branch back into an arm, or listing the scheme, breaks this test. + assert!(factory_result("oss://bucket/path", AccessMode::Read).is_ok()); + assert!(!builtin_storage_schemes(AccessMode::Write).contains(&"oss")); + let err = factory_result("oss://bucket/path", AccessMode::Write).unwrap_err(); + assert!( + err.contains("Unsupported storage scheme: oss"), + "unexpected error: {err}" + ); + assert!(factory_result("memory:///path", AccessMode::Write).is_ok()); + assert!(!builtin_storage_schemes(AccessMode::Read).contains(&"memory")); + let err = factory_result("memory:///path", AccessMode::Read).unwrap_err(); + assert!( + err.contains("Unsupported storage scheme: memory"), + "unexpected error: {err}" + ); + } + + #[test] + fn mixed_case_scheme_is_rejected_for_both_modes() { + // OpenDAL strips the scheme prefix from every path case-sensitively at open time + // (`S3://bucket/key` fails its `s3://bucket/` prefix check), so the factory must not + // admit what the open cannot serve. The JVM gates match the built-in set verbatim too. + for mode in [AccessMode::Read, AccessMode::Write] { + for path in ["S3://bucket/key", "File:///tmp/x", "GS://bucket/key"] { + let err = factory_result(path, mode).unwrap_err(); + assert!( + err.contains("Unsupported storage scheme"), + "unexpected error for {path} in {mode:?}: {err}" + ); + } + } + } + + #[test] + fn alias_scheme_is_matched_verbatim() { + // An alias location is opened as written, and OpenDAL's S3 backend checks it against a + // lowercase `scheme://bucket/` prefix, so a mixed-case alias that passed a + // case-insensitive gate here would only fail at open time. The list entry may be written + // in any case; the location's scheme must match its lowercase form exactly. + for listed in ["blob", " BLOB ", "minio,Blob"] { + let props = HashMap::from([( + "fs.comet.s3Compliant.schemes".to_string(), + listed.to_string(), + )]); + for mode in [AccessMode::Read, AccessMode::Write] { + assert!( + storage_factory_for("blob://bucket/key", &props, "test_cat", mode).is_ok(), + "blob://bucket/key must be admitted for {mode:?} with list {listed:?}" + ); + for path in [ + "BLOB://bucket/key", + "Blob://bucket/key", + "BLOB:///bucket/key", + ] { + let err = storage_factory_for(path, &props, "test_cat", mode) + .map(|_| ()) + .expect_err(&format!( + "{path} must be rejected for {mode:?} with list {listed:?}" + )) + .to_string(); + assert!( + err.contains("Unsupported storage scheme"), + "unexpected error for {path} in {mode:?}: {err}" + ); + } + } + } + } + #[test] fn unknown_scheme_is_rejected() { let err = factory_result("hdfs://nn/db/table", AccessMode::Read).unwrap_err(); @@ -869,6 +977,43 @@ mod tests { assert!(has_explicit_s3_credentials(&blank)); } + #[test] + fn listed_schemes_are_accepted_by_storage_factory() { + for mode in [AccessMode::Read, AccessMode::Write] { + for scheme in builtin_storage_schemes(mode) { + let url = format!("{scheme}://bucket/path"); + assert!( + factory_result(&url, mode).is_ok(), + "{scheme} is listed for {mode:?} but the factory rejects it" + ); + } + } + assert!(builtin_storage_schemes(AccessMode::Read).contains(&"oss")); + assert!(!builtin_storage_schemes(AccessMode::Write).contains(&"oss")); + assert!(!builtin_storage_schemes(AccessMode::Read).contains(&"memory")); + assert!(builtin_storage_schemes(AccessMode::Write).contains(&"memory")); + } + + #[test] + fn unlisted_schemes_are_rejected_for_both_modes() { + let unlisted = [ + "hdfs", "abfs", "abfss", "wasb", "wasbs", "gcs", "http", "https", "azure", + ]; + for mode in [AccessMode::Read, AccessMode::Write] { + for scheme in unlisted { + let err = factory_result(&format!("{scheme}://bucket/path"), mode).unwrap_err(); + assert!( + err.contains("Unsupported storage scheme"), + "unexpected error for {scheme} in {mode:?}: {err}" + ); + assert!( + !builtin_storage_schemes(mode).contains(&scheme), + "{scheme} must not be listed for {mode:?}" + ); + } + } + } + #[test] fn scheme_of_extracts_scheme_from_all_uri_forms() { // Host-bearing and hostless/opaque vendor forms must resolve to the same scheme, so an diff --git a/native/core/src/execution/operators/iceberg_scan.rs b/native/core/src/execution/operators/iceberg_scan.rs index 9f465329ca0..ad69cac4f03 100644 --- a/native/core/src/execution/operators/iceberg_scan.rs +++ b/native/core/src/execution/operators/iceberg_scan.rs @@ -230,7 +230,7 @@ impl IcebergScanExec { let scan_metrics = scan_result.metrics().clone(); let stream = scan_result.stream(); - let spark_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let spark_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); let adapter_factory = SparkPhysicalExprAdapterFactory::new(spark_options, None); let adapted_stream = @@ -630,7 +630,7 @@ mod tests { Field::new("a", DataType::Int64, false), Field::new("b", DataType::Int64, false), ])); - let mut options = super::SparkParquetOptions::new(super::EvalMode::Legacy, "UTC", false); + let mut options = super::SparkParquetOptions::new(super::EvalMode::Legacy, "UTC"); options.case_sensitive = true; let factory = super::SparkPhysicalExprAdapterFactory::new(options, None); for name in ["a", "b"] { diff --git a/native/core/src/execution/operators/iceberg_write.rs b/native/core/src/execution/operators/iceberg_write.rs index f7d5312fd92..ba9b08bd9da 100644 --- a/native/core/src/execution/operators/iceberg_write.rs +++ b/native/core/src/execution/operators/iceberg_write.rs @@ -27,6 +27,7 @@ use std::collections::HashMap; use std::fmt; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use arrow::array::{ @@ -35,8 +36,10 @@ use arrow::array::{ }; use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::datatypes::{DataType, Field, Schema as ArrowSchema, SchemaRef}; +use bytes::Bytes; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::error::{DataFusionError, Result as DFResult}; +use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion::execution::TaskContext; use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; @@ -48,12 +51,13 @@ use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, Partitioning, PlanProperties, SendableRecordBatchStream, }; +use futures::stream::BoxStream; use futures::TryStreamExt; use iceberg::arrow::arrow_struct_to_literal; -use iceberg::io::FileIO; +use iceberg::io::{FileIO, FileMetadata, FileRead, FileWrite, InputFile, OutputFile, Storage}; use iceberg::spec::{ - DataFile, DataFileFormat, Literal, ManifestWriterBuilder, PartitionKey, PartitionSpec, - PartitionSpecRef, Schema as IcebergSchema, SchemaRef as IcebergSchemaRef, + DataFile, DataFileBuilder, DataFileFormat, Literal, ManifestWriterBuilder, PartitionKey, + PartitionSpec, PartitionSpecRef, Schema as IcebergSchema, SchemaRef as IcebergSchemaRef, Struct as IcebergStruct, StructType, }; use iceberg::writer::base_writer::data_file_writer::{DataFileWriter, DataFileWriterBuilder}; @@ -61,16 +65,19 @@ use iceberg::writer::file_writer::location_generator::{ DefaultFileNameGenerator, LocationGenerator, }; use iceberg::writer::file_writer::rolling_writer::RollingFileWriterBuilder; -use iceberg::writer::file_writer::ParquetWriterBuilder; +use iceberg::writer::file_writer::{ + FileWriter, FileWriterBuilder, ParquetWriter, ParquetWriterBuilder, +}; use iceberg::writer::partitioning::clustered_writer::ClusteredWriter; use iceberg::writer::partitioning::fanout_writer::FanoutWriter; use iceberg::writer::partitioning::unpartitioned_writer::UnpartitionedWriter; -use iceberg::writer::IcebergWriterBuilder; +use iceberg::writer::{CurrentFileStatus, IcebergWriterBuilder}; use iceberg::ErrorKind; #[cfg(test)] use parquet::arrow::PARQUET_FIELD_ID_META_KEY; use parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel}; use parquet::file::properties::{EnabledStatistics, WriterProperties}; +use serde::{Deserialize, Serialize}; use datafusion_comet_proto::spark_operator::{ CompressionCodec as ProtoCompressionCodec, IcebergParquetWriteSettings, IcebergWrite, @@ -79,16 +86,20 @@ use datafusion_comet_proto::spark_operator::{ use crate::cloud::s3::credential_bridge::AccessMode; use crate::errors::CometError; -use crate::execution::operators::iceberg_common::load_file_io; +use crate::execution::operators::iceberg_common::{load_file_io, scheme_of}; use crate::execution::operators::iceberg_dictionary::DictionaryChooser; use crate::execution::operators::iceberg_partition_path::{ partition_to_path, CometLocationGenerator, }; use crate::execution::operators::iceberg_partition_value::PartitionValueCalculator; -/// The writer every partition gets: iceberg-rust's own data file writer, rolling Parquet files. -type PartitionDataFileWriter = - DataFileWriter; +/// The writer every partition gets: iceberg-rust's own data file writer, rolling Parquet files +/// that report what they hold in memory. +type PartitionDataFileWriter = DataFileWriter< + MeteredParquetWriterBuilder, + TrackingLocationGenerator, + DefaultFileNameGenerator, +>; /// Opens each partition's data file writer with the Parquet writer properties chosen for that /// partition. @@ -99,7 +110,8 @@ type PartitionDataFileWriter = /// `ParquetWriterBuilder` -> `RollingFileWriterBuilder` -> `DataFileWriterBuilder` chain once per /// partition, with that partition's properties. Every file the partition rolls into keeps them. /// The location and file name generators are shared, so file names still count up across the -/// task and every location is still tracked for cleanup. +/// task and every location is still tracked for cleanup, and the files all report to one +/// [`OpenFileMemory`]. struct PartitionWriterBuilder { schema: IcebergSchemaRef, target_file_size: usize, @@ -107,6 +119,8 @@ struct PartitionWriterBuilder { location_generator: TrackingLocationGenerator, file_name_generator: DefaultFileNameGenerator, properties: Arc, + open_files: OpenFileMemory, + storage: StorageWrites, } #[async_trait::async_trait] @@ -117,8 +131,13 @@ impl IcebergWriterBuilder for PartitionWriterBuilder { let properties = self .properties .take(partition_key.as_ref().map(PartitionKey::data)); + let parquet_builder = MeteredParquetWriterBuilder { + inner: ParquetWriterBuilder::new(properties, Arc::clone(&self.schema)), + open_files: self.open_files.clone(), + storage: self.storage, + }; let rolling_builder = RollingFileWriterBuilder::new( - ParquetWriterBuilder::new(properties, Arc::clone(&self.schema)), + parquet_builder, self.target_file_size, self.file_io.clone(), self.location_generator.clone(), @@ -175,6 +194,335 @@ impl PartitionProperties { } } +/// What a task's open data files hold in memory between them. +/// +/// iceberg-rust keeps each open file writer private inside its rolling and partitioning writers, +/// so the files report their own shares here through [`MeteredParquetWriter`], and `run_write_task` +/// reserves the total. A fanout write keeps one file open per partition, so this is what grows +/// with the partition count. +#[derive(Clone, Debug, Default)] +struct OpenFileMemory(Arc); + +impl OpenFileMemory { + fn bytes(&self) -> usize { + self.0.load(Ordering::Relaxed) + } + + fn share(&self) -> OpenFileShare { + OpenFileShare { + total: self.clone(), + bytes: 0, + } + } +} + +/// One open file's part of [`OpenFileMemory`], given back when the file closes, or when it is +/// dropped because the task failed mid-write. +#[derive(Debug)] +struct OpenFileShare { + total: OpenFileMemory, + bytes: usize, +} + +impl OpenFileShare { + fn set(&mut self, bytes: usize) { + if bytes > self.bytes { + self.total + .0 + .fetch_add(bytes - self.bytes, Ordering::Relaxed); + } else { + self.total + .0 + .fetch_sub(self.bytes - bytes, Ordering::Relaxed); + } + self.bytes = bytes; + } +} + +impl Drop for OpenFileShare { + fn drop(&mut self) { + self.set(0); + } +} + +/// [`ParquetWriterBuilder`] whose files report what they hold in memory to the task's +/// [`OpenFileMemory`]. +#[derive(Clone, Debug)] +struct MeteredParquetWriterBuilder { + inner: ParquetWriterBuilder, + open_files: OpenFileMemory, + /// How the files' storage takes the row groups they flush. + storage: StorageWrites, +} + +impl FileWriterBuilder for MeteredParquetWriterBuilder { + type R = MeteredParquetWriter; + + async fn build(&self, output_file: OutputFile) -> iceberg::Result { + let released = Arc::new(AtomicUsize::new(0)); + let output_file = CountedOutput::wrap(output_file, self.storage, Arc::clone(&released)); + Ok(MeteredParquetWriter { + inner: self.inner.build(output_file).await?, + share: self.open_files.share(), + released, + }) + } +} + +/// iceberg-rust's [`ParquetWriter`], reporting after every write how much of its file it still +/// holds in memory. +/// +/// parquet-rs keeps a file's in-progress row group in memory and hands it to storage once it +/// reaches the row-group size. `ParquetWriter` does not expose parquet-rs's estimate of that +/// memory, only `current_written_size`: every byte the file has written, the in-progress row +/// group's encoded size included. So the file's writer counts the bytes that leave memory on +/// their way to storage (see [`CountedOutput`]), and the share is what the file has written less +/// those: its in-progress row group, the few KiB parquet-rs buffers in front of storage, and what +/// the storage still holds of the row groups the file has flushed (see [`StorageWrites`]). +/// +/// The share does not see what parquet-rs keeps beyond the encoded estimate: dictionary +/// encoders' hash tables and unencoded indices, and buffer capacity past what is used. Nor would +/// it see Bloom filters, which parquet-rs sizes when a file opens and the eligibility gate +/// declines today. +struct MeteredParquetWriter { + inner: ParquetWriter, + share: OpenFileShare, + /// The bytes of the file that have left memory, counted by its [`CountingFileWrite`]. + released: Arc, +} + +impl FileWriter for MeteredParquetWriter { + async fn write(&mut self, batch: &RecordBatch) -> iceberg::Result<()> { + self.inner.write(batch).await?; + let released = self.released.load(Ordering::Relaxed); + self.share + .set(self.inner.current_written_size().saturating_sub(released)); + Ok(()) + } + + async fn close(self) -> iceberg::Result> { + // Closing flushes the last row group, and the file's buffers go with the writer. The share + // is given back when it drops at the end of this call, whether or not the close succeeded. + let Self { + inner, + share: _share, + .. + } = self; + inner.close().await + } +} + +impl CurrentFileStatus for MeteredParquetWriter { + fn current_file_path(&self) -> String { + self.inner.current_file_path() + } + + fn current_row_num(&self) -> usize { + self.inner.current_row_num() + } + + fn current_written_size(&self) -> usize { + self.inner.current_written_size() + } +} + +/// How the storage behind a task's data files takes the row groups a file flushes, which decides +/// when they leave memory. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum StorageWrites { + /// A local file is written as it goes, so a flushed row group leaves memory at once. + Through, + /// S3 and GCS take a file in parts. OpenDAL collects what a file flushes until it makes a part + /// of at least [`MIN_PART_BYTES`], and holds each part until the next one is complete or the + /// file closes, because only then does it know whether the part is the last. So a file whose + /// row groups are at least a part each keeps the last one it flushed in memory. + InParts, +} + +/// The smallest part OpenDAL uploads, its `write_multi_min_size` for both S3 and GCS. +const MIN_PART_BYTES: usize = 5 * 1024 * 1024; + +impl StorageWrites { + /// How the storage `load_file_io` builds for `location` takes a file. Native writes support + /// local storage (`file`, `memory`), S3 and GCS, and any other scheme is taken to upload in + /// parts, the way that reserves more. + fn for_location(location: &str) -> Self { + match scheme_of(location) { + "file" | "memory" => StorageWrites::Through, + _ => StorageWrites::InParts, + } + } +} + +/// What a file's storage holds of the bytes the file has handed it, following [`StorageWrites`]. +#[derive(Debug)] +struct HeldByStorage { + writes: StorageWrites, + /// Handed over since the last part was complete. + collecting: usize, + /// The last complete part, held until the next one. + last_part: usize, +} + +impl HeldByStorage { + fn new(writes: StorageWrites) -> Self { + Self { + writes, + collecting: 0, + last_part: 0, + } + } + + /// Takes `len` more bytes and returns how many bytes leave memory as it does: for a local + /// file these bytes, and for a file taken in parts the part held until now, if these complete + /// the next one. + fn hand_on(&mut self, len: usize) -> usize { + match self.writes { + StorageWrites::Through => len, + StorageWrites::InParts => { + self.collecting += len; + if self.collecting < MIN_PART_BYTES { + return 0; + } + std::mem::replace(&mut self.last_part, std::mem::take(&mut self.collecting)) + } + } + } +} + +/// The [`Storage`] behind one data file's [`OutputFile`], which opens the file's writer as a +/// [`CountingFileWrite`]. +/// +/// `ParquetWriter` opens its file's writer through the `OutputFile` it is built with, and an +/// `OutputFile` can only be made over a `Storage`. This one serves the one file it was made for, +/// through the `OutputFile` iceberg-rust made for it, and refuses what `ParquetWriter` never does +/// with its file. Serde exists only to satisfy the `Storage` typetag supertraits; Comet never +/// serializes it. +#[derive(Debug, Default, Serialize, Deserialize)] +struct CountedOutput { + #[serde(skip)] + file: Option, +} + +#[derive(Debug)] +struct CountedFile { + output: OutputFile, + writes: StorageWrites, + released: Arc, +} + +impl CountedOutput { + /// `output`, its writer counting into `released` the bytes that leave memory. + fn wrap(output: OutputFile, writes: StorageWrites, released: Arc) -> OutputFile { + let path = output.location().to_string(); + let file = CountedFile { + output, + writes, + released, + }; + OutputFile::new(Arc::new(CountedOutput { file: Some(file) }), path) + } + + fn file(&self, path: &str) -> iceberg::Result<&CountedFile> { + self.file + .as_ref() + .filter(|file| file.output.location() == path) + .ok_or_else(|| { + iceberg::Error::new( + ErrorKind::Unexpected, + format!("{path} is not the data file this output was made for"), + ) + }) + } + + fn unsupported(operation: &str, path: &str) -> iceberg::Error { + iceberg::Error::new( + ErrorKind::FeatureUnsupported, + format!("{operation} is not supported on the output of data file {path}"), + ) + } +} + +#[async_trait::async_trait] +#[typetag::serde(name = "CometCountedOutputStorage")] +impl Storage for CountedOutput { + async fn exists(&self, path: &str) -> iceberg::Result { + self.file(path)?.output.exists().await + } + + async fn metadata(&self, path: &str) -> iceberg::Result { + Err(Self::unsupported("metadata", path)) + } + + async fn read(&self, path: &str) -> iceberg::Result { + Err(Self::unsupported("read", path)) + } + + async fn reader(&self, path: &str) -> iceberg::Result> { + Err(Self::unsupported("reader", path)) + } + + async fn write(&self, path: &str, _bs: Bytes) -> iceberg::Result<()> { + Err(Self::unsupported("write", path)) + } + + async fn writer(&self, path: &str) -> iceberg::Result> { + let file = self.file(path)?; + Ok(Box::new(CountingFileWrite { + inner: file.output.writer().await?, + storage: HeldByStorage::new(file.writes), + released: Arc::clone(&file.released), + })) + } + + async fn delete(&self, path: &str) -> iceberg::Result<()> { + // `ParquetWriter` deletes a file it closes without having written a row. + self.file(path)?.output.delete().await + } + + async fn delete_prefix(&self, path: &str) -> iceberg::Result<()> { + Err(Self::unsupported("delete_prefix", path)) + } + + async fn delete_stream(&self, _paths: BoxStream<'static, String>) -> iceberg::Result<()> { + Err(Self::unsupported("delete_stream", "")) + } + + fn new_input(&self, path: &str) -> iceberg::Result { + Err(Self::unsupported("new_input", path)) + } + + fn new_output(&self, path: &str) -> iceberg::Result { + Err(Self::unsupported("new_output", path)) + } +} + +/// A data file's writer, counting into `released` the bytes that leave memory as the file hands +/// them to the storage's own writer, `inner`. +/// +/// parquet-rs hands a file's bytes on when it flushes a row group: everything it has written +/// since the last flush, except what is still in its write buffer. +struct CountingFileWrite { + inner: Box, + storage: HeldByStorage, + released: Arc, +} + +#[async_trait::async_trait] +impl FileWrite for CountingFileWrite { + async fn write(&mut self, bs: Bytes) -> iceberg::Result<()> { + let len = bs.len(); + self.inner.write(bs).await?; + self.released + .fetch_add(self.storage.hand_on(len), Ordering::Relaxed); + Ok(()) + } + + async fn close(&mut self) -> iceberg::Result { + self.inner.close().await + } +} + /// [`CometLocationGenerator`] that records every location it hands to a file writer. /// /// iceberg-rust's writers keep the `DataFile`s they have finalized private until `close`, and @@ -471,6 +819,10 @@ impl ExecutionPlan for IcebergWriteExec { // Time spent inside the iceberg-rust writer stack (write + close), excluding time spent // waiting on the upstream input stream. Surfaced on the JVM exec's SQL metrics by name. let write_time = MetricBuilder::new(&self.metrics).subset_time("write_time", partition); + // One consumer for the whole task, however many files it opens: a consumer per file + // would shrink every other consumer's share of the fair pool as a fanout write widened. + let reservation = MemoryConsumer::new(format!("IcebergWriteExec[{partition}]")) + .register(context.memory_pool()); let input_stream = self.input.execute(partition, context)?; let common = Arc::clone(&self.common); let iceberg_schema = Arc::clone(&self.iceberg_schema); @@ -492,6 +844,7 @@ impl ExecutionPlan for IcebergWriteExec { partition_id, task_attempt_id, write_time, + reservation, ) .await?; // The guard is still armed: until the output batch reaches the JVM nothing else knows @@ -561,6 +914,10 @@ impl DisplayAs for IcebergWriteExec { /// Iceberg field IDs, and routes through `UnpartitionedWriter`/`FanoutWriter`/`ClusteredWriter` /// depending on `writer_mode`. /// +/// After every batch, `reservation` is resized to what the task's writers hold between batches: +/// the open files' shares (see [`MeteredParquetWriter`]) and the rows waiting in their +/// [`PartitionFeed`]s. The writer cannot spill, so a reservation the pool refuses fails the task. +/// /// On success the still-armed [`AbortOnDrop`] is returned along with the data files: the caller /// owns cleanup until the JVM acknowledges the output after recording its locations. #[allow(clippy::too_many_arguments)] @@ -574,6 +931,7 @@ async fn run_write_task( partition_id: Option, task_attempt_id: Option, write_time: Time, + reservation: MemoryReservation, ) -> DFResult<(Vec, AbortOnDrop)> { // The JVM exec wrapper stamps both ids per task; a missing id means the plan template was // executed directly, and defaulting would make every task collide on the same file names. @@ -623,6 +981,7 @@ async fn run_write_task( writer_properties, &target_schema, )?)); + let open_files = OpenFileMemory::default(); let data_file_builder = PartitionWriterBuilder { schema: Arc::clone(&iceberg_schema), target_file_size: common.target_file_size_bytes as usize, @@ -630,6 +989,8 @@ async fn run_write_task( location_generator: location_generator.clone(), file_name_generator, properties: Arc::clone(&properties), + open_files: open_files.clone(), + storage: StorageWrites::for_location(&common.data_location), }; let mut abort_guard = AbortOnDrop { file_io, @@ -672,8 +1033,10 @@ async fn run_write_task( if nests_floats { decorated = drop_unwritten_values(decorated)?; } - let _timer = write_time.timer(); + let timer = write_time.timer(); writer.write(decorated, &properties).await?; + timer.done(); + reservation.try_resize(open_files.bytes() + writer.pending_bytes())?; } let _timer = write_time.timer(); writer.close(&properties).await @@ -718,6 +1081,23 @@ enum InnerWriter { } impl InnerWriter { + /// Memory held by the rows waiting in the feeds: those held back for a dictionary choice, and + /// up to `ROWS_DIVISOR - 1` rows for the rest of a unit, for every partition the task has + /// seen, in a fanout write. + fn pending_bytes(&self) -> usize { + match self { + InnerWriter::Unpartitioned(_, feed) => feed.pending_bytes(), + InnerWriter::Fanout(_, _, fanout) => fanout + .feeds + .values() + .map(|(_, feed)| feed.pending_bytes()) + .sum(), + InnerWriter::Clustered(_, _, live) => { + live.as_ref().map_or(0, |(_, feed)| feed.pending_bytes()) + } + } + } + /// Writes `batch` through the writer this task built, in the [`ROWS_DIVISOR`]-row units the /// rolling writer needs to re-check the target file size on iceberg-java's cadence. Rows are /// fed after partition splitting, so each partition's dictionary choice is made from its own @@ -1187,6 +1567,8 @@ struct RowPacer { /// Rows handed in but not yet handed over; fewer than [`ROWS_DIVISOR`] in total. pending: Vec, pending_rows: usize, + /// Memory held by `pending`. + pending_bytes: usize, } impl RowPacer { @@ -1195,6 +1577,7 @@ impl RowPacer { slicer, pending: Vec::new(), pending_rows: 0, + pending_bytes: 0, } } @@ -1217,6 +1600,7 @@ impl RowPacer { if offset < rows { let rest = self.slicer.detach(&batch, offset, rows - offset)?; self.pending_rows += rest.num_rows(); + self.pending_bytes += rest.get_array_memory_size(); self.pending.push(rest); } Ok(units) @@ -1243,6 +1627,7 @@ impl RowPacer { fn concat_pending(&mut self) -> DFResult { self.pending_rows = 0; + self.pending_bytes = 0; if self.pending.len() == 1 { return Ok(self.pending.pop().expect("pending has one batch")); } @@ -1295,6 +1680,12 @@ impl PartitionFeed { self.held.as_ref().map_or(0, |held| held.bytes) } + /// Memory taken by every row the feed has not handed to the writer: those held back for the + /// dictionary choice and those waiting in the pacer for the rest of their unit. + fn pending_bytes(&self) -> usize { + self.held_bytes() + self.pacer.pending_bytes + } + /// The complete units `batch` makes available, in row order. fn push( &mut self, @@ -2230,6 +2621,75 @@ mod tests { } } + /// The task reserves what the pacers hold back, so the count has to follow the rows: up while + /// they wait, and back to nothing once they are handed over. + #[test] + fn pacer_counts_the_memory_of_the_rows_it_holds_back() { + for gather in [false, true] { + let mut pacer = RowPacer::new(RowSlicer { gather }); + assert!(pacer.push(int_batch(800)).unwrap().is_empty()); + assert!( + pacer.pending_bytes >= 800 * std::mem::size_of::(), + "gather={gather}: {} bytes for 800 held rows", + pacer.pending_bytes + ); + // Completing the unit hands every held row over. + assert_eq!(pacer.push(int_batch_from(800, 200)).unwrap().len(), 1); + assert_eq!(pacer.pending_bytes, 0, "gather={gather}"); + assert!(pacer.push(int_batch(300)).unwrap().is_empty()); + assert!(pacer.pending_bytes > 0, "gather={gather}"); + assert!(pacer.flush().unwrap().is_some()); + assert_eq!(pacer.pending_bytes, 0, "gather={gather}"); + } + } + + /// A local file keeps nothing it hands on. S3 and GCS take a file in parts: OpenDAL collects + /// what it is handed until it makes a part of at least `MIN_PART_BYTES`, and holds each part + /// until the next one is complete, so a part leaves memory only when the next one is made. + #[test] + fn storage_taken_in_parts_holds_the_last_part_until_the_next() { + let mib = 1024 * 1024; + let mut local = HeldByStorage::new(StorageWrites::Through); + assert_eq!(local.hand_on(mib), mib); + + let mut parts = HeldByStorage::new(StorageWrites::InParts); + // Small row groups are collected until they make a part, which is then held. + for _ in 0..5 { + assert_eq!(parts.hand_on(mib), 0); + } + // A row group of a part or more is a part of its own, and sends the one held before it. + assert_eq!(parts.hand_on(6 * mib), 5 * mib); + assert_eq!(parts.hand_on(mib), 0); + assert_eq!(parts.hand_on(4 * mib), 6 * mib); + assert_eq!(parts.hand_on(8 * mib), 5 * mib); + } + + #[test] + fn only_local_storage_writes_through() { + for location in [ + "file:///tmp/warehouse", + "/tmp/warehouse", + "memory:///warehouse", + ] { + assert_eq!( + StorageWrites::for_location(location), + StorageWrites::Through, + "{location}" + ); + } + for location in [ + "s3://bucket/warehouse", + "s3a://bucket/warehouse", + "gs://bucket/w", + ] { + assert_eq!( + StorageWrites::for_location(location), + StorageWrites::InParts, + "{location}" + ); + } + } + // -- Integration tests against the real iceberg-rust writer stack --------- mod integration { @@ -2237,15 +2697,18 @@ mod tests { use arrow::array::{BinaryArray, Int32Array, StringArray, TimestampMicrosecondArray}; use arrow::datatypes::TimeUnit; use datafusion::common::Result as DFResult; + use datafusion::execution::memory_pool::{MemoryPool, UnboundedMemoryPool}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion_comet_proto::spark_operator::{ CompressionCodec as ProtoCodec, IcebergParquetWriteSettings, IcebergWriteCommon, IcebergWriterMode as ProtoIcebergWriterMode, }; use futures::StreamExt; + use iceberg::io::FileIOBuilder; use iceberg::spec::{ Manifest, NestedField, PartitionSpec, PrimitiveType, Schema, Transform, Type, }; + use iceberg_storage_opendal::OpenDalStorageFactory; use parquet::file::properties::WriterProperties; use std::collections::HashMap; use std::path::PathBuf; @@ -2253,6 +2716,8 @@ mod tests { use std::time::Duration; use tempfile::TempDir; + use crate::execution::memory_pools::testing::PeakMemoryPool; + fn user_schema() -> SchemaRef { Arc::new(ArrowSchema::new(vec![ Field::new("id", DataType::Int32, false), @@ -2296,6 +2761,12 @@ mod tests { )) } + /// A reservation no write can outgrow, for the tests that are not about memory. + fn unbounded_reservation() -> MemoryReservation { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + MemoryConsumer::new("IcebergWriteExec[0]").register(&pool) + } + fn common( data_location: String, spec_json: String, @@ -2362,6 +2833,7 @@ mod tests { Some(0), Some(0), Time::default(), + unbounded_reservation(), ) .await?; // These tests assert on the written files, so they stand in for the JVM taking @@ -2461,6 +2933,7 @@ mod tests { Some(0), Some(0), Time::default(), + unbounded_reservation(), ) .await .unwrap(); @@ -2508,6 +2981,7 @@ mod tests { Some(0), Some(0), Time::default(), + unbounded_reservation(), ) .await .unwrap(); @@ -3784,6 +4258,7 @@ mod tests { Some(0), Some(0), Time::default(), + unbounded_reservation(), ) .await .unwrap(); @@ -3912,6 +4387,7 @@ mod tests { Some(0), Some(0), Time::default(), + unbounded_reservation(), ) .await .unwrap(); @@ -3979,6 +4455,7 @@ mod tests { Some(0), Some(0), Time::default(), + unbounded_reservation(), ) .await .unwrap(); @@ -4071,6 +4548,7 @@ mod tests { Some(0), Some(0), Time::default(), + unbounded_reservation(), ) .await .unwrap(); @@ -4087,6 +4565,475 @@ mod tests { ); } + // -- Memory accounting --------------------------------------------------------------- + + const TARGET_FILE_SIZE: u64 = 512 * 1024 * 1024; + + /// parquet-rs writes a file through a `BufWriter` of the default capacity, so up to this + /// much of a row group it has flushed can still be in memory after the flush. + const PARQUET_WRITE_BUFFER_BYTES: usize = 8 * 1024; + + /// Rows `first..first + rows`, spread round robin over the regions `r0`..`r{partitions}`. + fn round_robin_batch_from(first: usize, rows: usize, partitions: usize) -> RecordBatch { + let ids: Vec = (first..first + rows).map(|id| id as i32).collect(); + let regions: Vec = (0..rows) + .map(|row| format!("r{}", row % partitions)) + .collect(); + let regions: Vec<&str> = regions.iter().map(String::as_str).collect(); + batch(&ids, ®ions) + } + + fn round_robin_batch(rows: usize, partitions: usize) -> RecordBatch { + round_robin_batch_from(0, rows, partitions) + } + + /// As many rows for each of the regions `r0`..`r{partitions}`, one region after another, + /// the way a clustered write receives them. + fn clustered_batch(rows: usize, partitions: usize) -> RecordBatch { + let ids: Vec = (0..rows as i32).collect(); + let regions: Vec = (0..rows) + .map(|row| format!("r{}", row * partitions / rows)) + .collect(); + let regions: Vec<&str> = regions.iter().map(String::as_str).collect(); + batch(&ids, ®ions) + } + + /// [`iceberg_user_schema`] with a `payload` column, for rows large enough to fill row + /// groups well past what parquet-rs buffers in front of storage. + fn iceberg_payload_schema() -> Schema { + Schema::builder() + .with_schema_id(1) + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::required(2, "region", Type::Primitive(PrimitiveType::String)) + .into(), + NestedField::required(3, "payload", Type::Primitive(PrimitiveType::String)) + .into(), + ]) + .build() + .unwrap() + } + + /// Rows `first..first + rows` of `region`, each with its own `payload_bytes`-byte payload. + fn payload_batch_from( + region: &str, + first: usize, + rows: usize, + payload_bytes: usize, + ) -> RecordBatch { + let ids: Vec = (first..first + rows).map(|id| id as i32).collect(); + let payloads: Vec = (first..first + rows) + .map(|id| format!("{id:0>payload_bytes$}")) + .collect(); + RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("region", DataType::Utf8, false), + Field::new("payload", DataType::Utf8, false), + ])), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(StringArray::from(vec![region; rows])), + Arc::new(StringArray::from(payloads)), + ], + ) + .unwrap() + } + + /// Writes `batches` into a fresh table of `schema`, partitioned by `region` unless the + /// mode is unpartitioned, reserving from `pool` the way `IcebergWriteExec::execute` + /// reserves from the task's pool. A successful write's files are left for the caller to + /// inspect. + async fn write_reserving_from( + pool: &Arc, + schema: Schema, + writer_mode: ProtoIcebergWriterMode, + batches: Vec, + writer_properties: WriterProperties, + target_file_size_bytes: u64, + ) -> (TempDir, DFResult>) { + let temp_dir = TempDir::new().unwrap(); + let spec = match writer_mode { + ProtoIcebergWriterMode::IcebergWriterUnpartitioned => { + PartitionSpec::builder(Arc::new(schema.clone())) + .build() + .unwrap() + } + _ => identity_region_spec(&schema), + }; + let common = with_target_file_size( + common( + format!("file://{}", temp_dir.path().display()), + serde_json::to_string(&spec).unwrap(), + serde_json::to_string(&schema).unwrap(), + writer_mode, + ), + target_file_size_bytes, + ); + let pool: Arc = Arc::::clone(pool); + let reservation = MemoryConsumer::new("IcebergWriteExec[0]").register(&pool); + let written = run_write_task( + input_stream(batches), + common, + Arc::new(schema), + Arc::new(spec), + writer_mode, + writer_properties, + Some(0), + Some(0), + Time::default(), + reservation, + ) + .await + .map(|(data_files, mut abort_guard)| { + abort_guard.disarm(); + data_files + }); + (temp_dir, written) + } + + /// How many files one batch was written to, and the most the write had reserved. Every + /// write has to give back what it reserved. + async fn peak_reservation( + writer_mode: ProtoIcebergWriterMode, + batch: RecordBatch, + target_file_size_bytes: u64, + ) -> (usize, usize) { + let pool = Arc::new(PeakMemoryPool::new(usize::MAX)); + // A partition's first page of rows is held back for its dictionary choice. One unit + // is that page here, so every unit goes straight to its file. + let properties = WriterProperties::builder() + .set_data_page_row_count_limit(ROWS_DIVISOR) + .build(); + let (_dir, written) = write_reserving_from( + &pool, + iceberg_user_schema(), + writer_mode, + vec![batch], + properties, + target_file_size_bytes, + ) + .await; + let files = written.unwrap().len(); + assert_eq!(pool.reserved(), 0, "{writer_mode:?} kept a reservation"); + (files, pool.peak()) + } + + fn parquet_files_under(dir: &std::path::Path) -> Vec { + let mut found = Vec::new(); + let mut dirs = vec![dir.to_path_buf()]; + while let Some(dir) = dirs.pop() { + for entry in std::fs::read_dir(dir).unwrap() { + let path = entry.unwrap().path(); + if path.is_dir() { + dirs.push(path); + } else if path.extension().is_some_and(|ext| ext == "parquet") { + found.push(path); + } + } + } + found + } + + /// A fanout partition holds rows back until they fill a unit, and every partition value + /// has its own. Here no partition fills one, so nothing reaches a file before the close: + /// the reservation is the held-back rows alone, and it grows with the partitions holding + /// them. + #[tokio::test] + async fn a_fanout_write_reserves_the_rows_its_partitions_hold_back() { + let rows_per_partition = 900; + let mut peaks = Vec::new(); + for partitions in [1, 8] { + let pool = Arc::new(PeakMemoryPool::new(usize::MAX)); + let (_dir, written) = write_reserving_from( + &pool, + iceberg_user_schema(), + ProtoIcebergWriterMode::IcebergWriterFanout, + vec![round_robin_batch( + rows_per_partition * partitions, + partitions, + )], + WriterProperties::builder().build(), + TARGET_FILE_SIZE, + ) + .await; + assert_eq!( + record_counts(&written.unwrap()), + vec![rows_per_partition as u64; partitions] + ); + assert_eq!(pool.reserved(), 0, "the write kept a reservation"); + peaks.push(pool.peak()); + } + assert!(peaks[0] > 0, "held-back rows were not reserved: {peaks:?}"); + assert!( + peaks[1] > 4 * peaks[0], + "the reservation did not grow with the partitions holding rows back: {peaks:?}" + ); + } + + /// An open file reserves what it holds and gives it back when it closes. A fanout write + /// keeps every partition's file open, a clustered write one partition's at a time, and a + /// rolling writer one file however many it rolls through. + #[tokio::test] + async fn the_reservation_follows_the_files_that_are_open() { + // Eight partitions of exactly one unit each: no rows are waiting for a full unit, so + // after the batch every row is in a file. + let rows = 8 * ROWS_DIVISOR; + let (fanout_files, fanout) = peak_reservation( + ProtoIcebergWriterMode::IcebergWriterFanout, + round_robin_batch(rows, 8), + TARGET_FILE_SIZE, + ) + .await; + let (clustered_files, clustered) = peak_reservation( + ProtoIcebergWriterMode::IcebergWriterClustered, + clustered_batch(rows, 8), + TARGET_FILE_SIZE, + ) + .await; + // A one-byte target rolls to a new file for every unit. + let (rolled_files, rolled) = peak_reservation( + ProtoIcebergWriterMode::IcebergWriterUnpartitioned, + round_robin_batch(rows, 8), + 1, + ) + .await; + assert_eq!((fanout_files, clustered_files, rolled_files), (8, 8, 8)); + assert!( + clustered > 0 && rolled > 0, + "open files were not reserved: clustered {clustered}, rolled {rolled}" + ); + // One file open against eight. A closed file that kept its share would leave the + // clustered and rolled writes holding all eight too. + assert!( + 4 * clustered < fanout, + "clustered {clustered}, fanout {fanout}" + ); + assert!(4 * rolled < fanout, "rolled {rolled}, fanout {fanout}"); + } + + /// parquet-rs holds only a file's in-progress row group and what it buffers in front of + /// storage; a local file keeps none of the row groups it has flushed. A file written + /// through many row groups reserves no more than one and that buffer, although its written + /// size keeps growing. + #[tokio::test] + async fn a_file_reserves_no_more_than_one_row_group() { + let row_group_bytes = 16 * 1024; + let properties = WriterProperties::builder() + .set_max_row_group_bytes(Some(row_group_bytes)) + .build(); + // Whole units, so nothing is held back between batches. + let batches = (0..50) + .map(|unit| round_robin_batch_from(unit * ROWS_DIVISOR, ROWS_DIVISOR, 8)) + .collect(); + let pool = Arc::new(PeakMemoryPool::new(usize::MAX)); + let (_dir, written) = write_reserving_from( + &pool, + iceberg_user_schema(), + ProtoIcebergWriterMode::IcebergWriterUnpartitioned, + batches, + properties, + TARGET_FILE_SIZE, + ) + .await; + let data_files = written.unwrap(); + assert_eq!(data_files.len(), 1); + let file_size = data_files[0].file_size_in_bytes() as usize; + assert!( + file_size > 8 * row_group_bytes, + "the file never outgrew a row group: {file_size} bytes" + ); + assert!(pool.peak() > 0, "the open file was not reserved"); + assert!( + pool.peak() <= row_group_bytes + PARQUET_WRITE_BUFFER_BYTES, + "reserved {} bytes for a {row_group_bytes}-byte row group", + pool.peak() + ); + } + + /// A local file keeps nothing of the row groups it has flushed, so a file that flushes + /// gives back what it reserved for them. Here every partition's unit is larger than a row + /// group and flushes as soon as it is written, so twelve open files together hold less + /// than one row group, and a pool of eight row groups holds the write. Charging a flushed + /// file a whole row group until it closed failed this write on its ninth partition. + #[tokio::test] + async fn files_that_flush_give_back_what_they_reserved() { + let row_group_bytes = 128 * 1024; + let properties = WriterProperties::builder() + // One unit is one page, so no rows wait for the dictionary choice. + .set_data_page_row_count_limit(ROWS_DIVISOR) + .set_max_row_group_bytes(Some(row_group_bytes)) + .build(); + // One batch per partition, each a unit of distinct 256-byte payloads: about 256 KB. + let batches = (0..12) + .map(|partition| payload_batch_from(&format!("r{partition}"), 0, ROWS_DIVISOR, 256)) + .collect(); + let pool = Arc::new(PeakMemoryPool::new(8 * row_group_bytes)); + let (_dir, written) = write_reserving_from( + &pool, + iceberg_payload_schema(), + ProtoIcebergWriterMode::IcebergWriterFanout, + batches, + properties, + TARGET_FILE_SIZE, + ) + .await; + let data_files = written.expect("flushed row groups stayed reserved"); + assert_eq!(data_files.len(), 12); + assert!( + data_files + .iter() + .all(|file| file.file_size_in_bytes() > row_group_bytes as u64), + "a partition's file did not outgrow a row group" + ); + assert!( + pool.peak() < row_group_bytes, + "twelve flushed files reserved {} bytes", + pool.peak() + ); + assert_eq!(pool.reserved(), 0, "the write kept a reservation"); + } + + /// A write whose open files outgrow what the pool grants fails with the pool's error, + /// rather than holding memory nothing accounts for, and deletes the files it had opened. + #[tokio::test] + async fn a_write_the_pool_cannot_hold_fails_and_deletes_its_files() { + let batches = || vec![round_robin_batch(16 * ROWS_DIVISOR, 16)]; + let properties = || WriterProperties::builder().build(); + + // With room to spare, every partition's unit reaches its own file. + let roomy = Arc::new(PeakMemoryPool::new(usize::MAX)); + let (dir, written) = write_reserving_from( + &roomy, + iceberg_user_schema(), + ProtoIcebergWriterMode::IcebergWriterFanout, + batches(), + properties(), + TARGET_FILE_SIZE, + ) + .await; + assert_eq!(written.unwrap().len(), 16); + assert_eq!(parquet_files_under(dir.path()).len(), 16); + + let tight = Arc::new(PeakMemoryPool::new(roomy.peak() / 2)); + let (dir, written) = write_reserving_from( + &tight, + iceberg_user_schema(), + ProtoIcebergWriterMode::IcebergWriterFanout, + batches(), + properties(), + TARGET_FILE_SIZE, + ) + .await; + let error = written.expect_err("a write that outgrows the pool must fail"); + assert!( + matches!(error, DataFusionError::ResourcesExhausted(_)), + "{error:?}" + ); + assert!(error.to_string().contains("IcebergWriteExec[0]"), "{error}"); + assert_eq!( + parquet_files_under(dir.path()), + Vec::::new(), + "the failed write left files behind" + ); + assert_eq!(tight.reserved(), 0, "the failed write kept a reservation"); + } + + /// A task that fails mid-write drops its open files without closing them, so a file's + /// share has to come back either way. + #[tokio::test] + async fn an_open_file_gives_back_its_share_when_it_closes_or_is_dropped() { + let schema = Arc::new(iceberg_user_schema()); + let target = Arc::new(iceberg::arrow::schema_to_arrow_schema(&schema).unwrap()); + let rows = + decorate_batch_with_field_ids(round_robin_batch(ROWS_DIVISOR, 1), &target).unwrap(); + let file_io = FileIOBuilder::new(Arc::new(OpenDalStorageFactory::Memory)).build(); + let open_files = OpenFileMemory::default(); + let builder = MeteredParquetWriterBuilder { + inner: ParquetWriterBuilder::new(WriterProperties::builder().build(), schema), + open_files: open_files.clone(), + storage: StorageWrites::Through, + }; + let mut closed = builder + .build(file_io.new_output("memory:/t/closed.parquet").unwrap()) + .await + .unwrap(); + let mut dropped = builder + .build(file_io.new_output("memory:/t/dropped.parquet").unwrap()) + .await + .unwrap(); + + closed.write(&rows).await.unwrap(); + let one = open_files.bytes(); + assert!(one > 0, "a written file reported nothing"); + dropped.write(&rows).await.unwrap(); + assert_eq!(open_files.bytes(), 2 * one); + + closed.close().await.unwrap(); + assert_eq!(open_files.bytes(), one); + drop(dropped); + assert_eq!(open_files.bytes(), 0); + } + + /// A file's share is what it has written and still holds. A local file keeps nothing it + /// has flushed, so its share falls back whenever a row group flushes and never passes a + /// row group and parquet-rs's write buffer. A file taken in parts keeps what it flushes + /// until that makes a part, so here, short of a part, it holds everything it has written. + #[tokio::test] + async fn a_files_share_is_what_it_still_holds() { + let row_group_bytes = 128 * 1024; + let schema = Arc::new(iceberg_payload_schema()); + let target = Arc::new(iceberg::arrow::schema_to_arrow_schema(&schema).unwrap()); + let file_io = FileIOBuilder::new(Arc::new(OpenDalStorageFactory::Memory)).build(); + for storage in [StorageWrites::Through, StorageWrites::InParts] { + let open_files = OpenFileMemory::default(); + let builder = MeteredParquetWriterBuilder { + inner: ParquetWriterBuilder::new( + WriterProperties::builder() + .set_max_row_group_bytes(Some(row_group_bytes)) + .build(), + Arc::clone(&schema), + ), + open_files: open_files.clone(), + storage, + }; + let location = format!("memory:/t/{storage:?}.parquet"); + let mut file = builder + .build(file_io.new_output(&location).unwrap()) + .await + .unwrap(); + let mut fell_back = false; + let mut previous = 0; + // Units of about 40 KB, so a row group flushes every three or four. + for unit in 0..20 { + let rows = payload_batch_from("r0", unit * ROWS_DIVISOR, ROWS_DIVISOR, 32); + let rows = decorate_batch_with_field_ids(rows, &target).unwrap(); + file.write(&rows).await.unwrap(); + let share = open_files.bytes(); + match storage { + StorageWrites::Through => assert!( + share <= row_group_bytes + PARQUET_WRITE_BUFFER_BYTES, + "unit {unit}: a local file reserved {share} bytes" + ), + StorageWrites::InParts => { + assert_eq!(share, file.current_written_size(), "unit {unit}") + } + } + fell_back |= share < previous; + previous = share; + } + assert!( + file.current_written_size() > 4 * row_group_bytes, + "{storage:?}: the file flushed too few row groups" + ); + if storage == StorageWrites::Through { + assert!(fell_back, "the share never fell back as row groups flushed"); + } + file.close().await.unwrap(); + assert_eq!(open_files.bytes(), 0, "{storage:?}"); + } + } + /// Regression for apache/datafusion-comet#6145. Past `chrono`'s calendar (year 262142) the /// `day` and `identity(date)` directory names used to panic the task while iceberg-rust /// rendered them, and `year` and `month` came back as NULL partition values. The expected diff --git a/native/core/src/execution/operators/parquet_writer.rs b/native/core/src/execution/operators/parquet_writer.rs index 9897ed626c6..cc2da616e54 100644 --- a/native/core/src/execution/operators/parquet_writer.rs +++ b/native/core/src/execution/operators/parquet_writer.rs @@ -39,7 +39,7 @@ use async_trait::async_trait; use datafusion::{ common::tree_node::TreeNodeRecursion, error::{DataFusionError, Result}, - execution::context::TaskContext, + execution::{context::TaskContext, memory_pool::MemoryConsumer}, physical_expr::{EquivalenceProperties, PhysicalExpr}, physical_plan::{ execution_plan::{Boundedness, EmissionType}, @@ -157,6 +157,20 @@ impl ParquetWriter { } } + /// Memory the writer holds between batches: parquet-rs's estimate for the in-progress row + /// group, which counts encoded pages, encoder buffers, dictionaries and any Bloom filters. + /// The remote writer flushes a row group after every batch and stages it in a buffer that is + /// cleared after each upload but keeps its capacity, so that buffer counts too. + fn memory_size(&self) -> usize { + match self { + ParquetWriter::LocalFile(writer) => writer.memory_size(), + #[cfg(feature = "hdfs-opendal")] + ParquetWriter::Remote(writer, ..) => { + writer.memory_size() + writer.inner().get_ref().capacity() + } + } + } + /// Close the writer and finalize the file async fn close(self) -> std::result::Result<(), parquet::errors::ParquetError> { match self { @@ -220,18 +234,9 @@ impl ParquetWriter { pub struct ParquetWriterExec { /// Input execution plan input: Arc, - /// Where this task writes. When `work_dir` is set (the Spark 3.x `CometNativeWriteExec` - /// path) this is the write's output directory and is unused; the file name is derived from - /// `work_dir`. Otherwise (Spark 4.0+, `CometWriteFilesExec`) it is the exact path of the file - /// to write, chosen by the JVM commit protocol and used verbatim - this operator then never - /// derives file names of its own. + /// The exact path of the file this task writes, chosen by Spark's commit protocol on the JVM + /// side and used verbatim - this operator never derives file names of its own. output_path: String, - /// Working directory for temporary files (used by FileCommitProtocol). Spark 3.x only. - work_dir: Option, - /// Job ID for tracking this write operation - job_id: Option, - /// Task attempt ID for this specific task - task_attempt_id: Option, /// Compression codec compression: ParquetCompression, /// Partition ID (from Spark TaskContext) @@ -254,9 +259,6 @@ impl ParquetWriterExec { pub fn try_new( input: Arc, output_path: String, - work_dir: Option, - job_id: Option, - task_attempt_id: Option, compression: ParquetCompression, partition_id: i32, column_names: Vec, @@ -276,9 +278,6 @@ impl ParquetWriterExec { Ok(ParquetWriterExec { input, output_path, - work_dir, - job_id, - task_attempt_id, compression, partition_id, column_names, @@ -464,9 +463,6 @@ impl ExecutionPlan for ParquetWriterExec { 1 => Ok(Arc::new(ParquetWriterExec::try_new( Arc::clone(&children[0]), self.output_path.clone(), - self.work_dir.clone(), - self.job_id.clone(), - self.task_attempt_id, self.compression.clone(), self.partition_id, self.column_names.clone(), @@ -492,10 +488,10 @@ impl ExecutionPlan for ParquetWriterExec { let rows_written = MetricBuilder::new(&self.metrics).counter("rows_written", partition); let runtime_env = context.runtime_env(); + let reservation = MemoryConsumer::new(format!("ParquetWriterExec[{partition}]")) + .register(context.memory_pool()); let input = self.input.execute(partition, context)?; let input_schema = self.input.schema(); - let work_dir = self.work_dir.clone(); - let task_attempt_id = self.task_attempt_id; let compression = self.compression.to_parquet()?; let column_names = self.column_names.clone(); @@ -514,19 +510,8 @@ impl ExecutionPlan for ParquetWriterExec { Arc::new(Schema::new(fields)) }); - let part_file = match &work_dir { - // Spark 4.0+ hands over the exact file to write, chosen by the JVM commit protocol. - None => self.output_path.clone(), - // Spark 3.x hands over a working directory instead and expects the writer to name the - // file; that branch goes away with Spark 3.x support. - Some(work_dir) => match task_attempt_id { - Some(attempt_id) => format!( - "{}/part-{:05}-{:05}.parquet", - work_dir, self.partition_id, attempt_id - ), - None => format!("{}/part-{:05}.parquet", work_dir, self.partition_id), - }, - }; + // The JVM commit protocol has already chosen the exact file to write. + let part_file = self.output_path.clone(); // Configure writer properties let props = WriterProperties::builder() @@ -577,6 +562,8 @@ impl ExecutionPlan for ParquetWriterExec { writer.write(&renamed_batch).await.map_err(|e| { DataFusionError::Execution(format!("Failed to write batch: {}", e)) })?; + // The writer cannot spill, so a row group the pool will not grant fails the task. + reservation.try_resize(writer.memory_size())?; } writer.close().await.map_err(|e| { @@ -619,11 +606,14 @@ impl ExecutionPlan for ParquetWriterExec { #[cfg(test)] mod tests { use super::*; + use crate::execution::memory_pools::testing::PeakMemoryPool; use arrow::array::{Array, Int32Array, ListArray, StringArray}; use arrow::datatypes::{DataType, Field, Int32Type, Schema}; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; - use datafusion::prelude::SessionContext; + use datafusion::execution::memory_pool::MemoryPool; + use datafusion::execution::runtime_env::RuntimeEnvBuilder; + use datafusion::prelude::{SessionConfig, SessionContext}; use parquet::arrow::PARQUET_FIELD_ID_META_KEY; use parquet::basic::Repetition; use parquet::file::reader::{FileReader, SerializedFileReader}; @@ -654,11 +644,10 @@ mod tests { ); } - /// Spark 4.0+ hands over the exact file to write rather than a working directory. The writer - /// must use that path verbatim - Spark's commit protocol owns naming and staging, and + /// The JVM hands over the exact file to write. The writer must use that path verbatim - Spark's commit protocol owns naming and staging, and /// committers that track individual files depend on the name it chose. #[tokio::test] - async fn test_parquet_writer_uses_output_path_verbatim_without_work_dir() -> Result<()> { + async fn test_parquet_writer_uses_output_path_verbatim() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, true)])); let batch = RecordBatch::try_new( Arc::clone(&schema), @@ -675,9 +664,6 @@ mod tests { let writer = ParquetWriterExec::try_new( input, output_path, - None, // work_dir: Spark 4.0+ path - None, - None, ParquetCompression::None, // A non-zero partition id must not leak into the file name. 3, @@ -742,13 +728,10 @@ mod tests { let memory_source = MemorySourceConfig::try_new(&[vec![batch]], input_schema, None)?; let input = Arc::new(DataSourceExec::new(Arc::new(memory_source))); let temp_dir = tempfile::tempdir()?; - let work_dir = format!("file://{}", temp_dir.path().display()); + let output_path = format!("file://{}/part-00000.parquet", temp_dir.path().display()); let writer = ParquetWriterExec::try_new( input, - work_dir.clone(), - Some(work_dir), - None, - None, + output_path, ParquetCompression::None, 0, vec!["required_id".to_string(), "values".to_string()], @@ -811,13 +794,10 @@ mod tests { let memory_source = MemorySourceConfig::try_new(&[vec![batch]], input_schema, None)?; let input = Arc::new(DataSourceExec::new(Arc::new(memory_source))); let temp_dir = tempfile::tempdir()?; - let work_dir = format!("file://{}", temp_dir.path().display()); + let output_path = format!("file://{}/part-00000.parquet", temp_dir.path().display()); let writer = ParquetWriterExec::try_new( input, - work_dir.clone(), - Some(work_dir), - None, - None, + output_path, ParquetCompression::None, 0, vec!["values".to_string()], @@ -845,7 +825,6 @@ mod tests { /// Helper function to create a test RecordBatch with 1000 rows of (int, string) data /// Example batch_id 1 -> 0..1000, 2 -> 1001..2000 - #[allow(dead_code)] fn create_test_record_batch(batch_id: i32) -> Result { assert!(batch_id > 0, "batch_id must be greater than 0"); let num_rows = batch_id * 1000; @@ -868,6 +847,75 @@ mod tests { .map_err(|e| DataFusionError::ArrowError(Box::new(e), None)) } + /// Writes `batches` to a local file, reserving from `pool` as a task reserves from its own, + /// and returns the size of the file written. + async fn write_reserving_from( + pool: &Arc, + batches: Vec, + ) -> Result { + let schema = batches[0].schema(); + let column_names = schema.fields().iter().map(|f| f.name().clone()).collect(); + let memory_source = MemorySourceConfig::try_new(&[batches], Arc::clone(&schema), None)?; + let input = Arc::new(DataSourceExec::new(Arc::new(memory_source))); + let temp_dir = tempfile::tempdir()?; + let file = temp_dir.path().join("part-00000.parquet"); + let writer = ParquetWriterExec::try_new( + input, + format!("file://{}", file.display()), + ParquetCompression::None, + 0, + column_names, + None, + HashMap::new(), + )?; + let pool: Arc = Arc::::clone(pool); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(pool) + .build_arc()?; + let context = SessionContext::new_with_config_rt(SessionConfig::new(), runtime); + let mut stream = writer.execute(0, context.task_ctx())?; + while stream.try_next().await?.is_some() {} + Ok(std::fs::metadata(file)?.len()) + } + + /// The writer reserves its in-progress row group, here the whole file, until the file is + /// closed. + #[tokio::test] + async fn test_parquet_writer_reserves_its_row_group_until_the_file_closes() -> Result<()> { + let pool = Arc::new(PeakMemoryPool::new(usize::MAX)); + let batches = (1..=5) + .map(create_test_record_batch) + .collect::>()?; + let file_size = write_reserving_from(&pool, batches).await?; + assert!( + pool.peak() as u64 >= file_size / 2, + "reserved {} bytes for a row group written as a {file_size}-byte file", + pool.peak() + ); + assert_eq!(pool.reserved(), 0, "the writer kept its reservation"); + Ok(()) + } + + /// A row group the pool will not grant fails the write with the pool's error, rather than + /// holding memory nothing accounts for. + #[tokio::test] + async fn test_parquet_writer_fails_when_the_pool_refuses_its_row_group() -> Result<()> { + let pool = Arc::new(PeakMemoryPool::new(1024)); + let error = write_reserving_from(&pool, vec![create_test_record_batch(1)?]) + .await + .expect_err("a row group larger than the pool must fail the write"); + assert!( + matches!(error, DataFusionError::ResourcesExhausted(_)), + "{error:?}" + ); + assert!( + error.to_string().contains("ParquetWriterExec[0]"), + "{error}" + ); + assert_eq!(pool.reserved(), 0, "the failed write kept its reservation"); + Ok(()) + } + #[tokio::test] #[cfg(feature = "hdfs-opendal")] #[ignore = "This test requires a running HDFS cluster"] @@ -1068,16 +1116,14 @@ mod tests { let memory_exec = Arc::new(DataSourceExec::new(Arc::new(memory_source_config))); // Create ParquetWriterExec with DataSourceExec as input - let output_path = "unused".to_string(); - let work_dir = "hdfs://namenode:9000/user/test_parquet_writer_exec".to_string(); + let output_path = + "hdfs://namenode:9000/user/test_parquet_writer_exec/part-00000-00123.parquet" + .to_string(); let column_names = vec!["id".to_string(), "name".to_string()]; let parquet_writer = ParquetWriterExec::try_new( memory_exec, output_path, - Some(work_dir), - None, // job_id - Some(123), // task_attempt_id ParquetCompression::None, 0, // partition_id column_names, diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 99a52fb562f..b17c7833bc3 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -39,6 +39,7 @@ use crate::execution::operators::DynamicFilterJoinExec; use crate::execution::operators::IcebergScanExec; use crate::execution::operators::TopKReaderFilterExec; use crate::execution::{ + expressions::map_entries::MapEntriesExpr, operators::{ ExecutionError, MergeActionContext, MergeInstructionExec, MergeRowsExec, ScanExec, ShuffleScanExec, @@ -144,10 +145,10 @@ use datafusion_comet_proto::{ use datafusion_comet_spark_expr::{ create_case_when, create_if_expr, jvm_udf::JvmScalarUdfExpr, spark_in_list, ApproxPercentile, ArrayInsert, Avg, AvgDecimal, Cast, CheckOverflow, Correlation, Covariance, CreateNamedStruct, - DecimalRescaleCheckOverflow, GetArrayStructFields, GetStructField, HllPlusPlus, HllSketchAgg, - HllUnionAgg, IfExpr, ListExtract, MaxMinBy, Mode, NormalizeNaNAndZero, NormalizeNestedFloats, - Regr, RegrType, SparkCastOptions, SparkMinMax, Stddev, SumDecimal, ToJson, UnboundColumn, - Variance, WideDecimalBinaryExpr, WideDecimalOp, + DecimalRescaleCheckOverflow, FloatOperands, GetArrayStructFields, GetStructField, HllPlusPlus, + HllSketchAgg, HllUnionAgg, IfExpr, ListExtract, MaxMinBy, Mode, NormalizeNaNAndZero, + NormalizeNestedFloats, Regr, RegrType, SparkCastOptions, SparkMinMax, Stddev, SumDecimal, + ToJson, UnboundColumn, Variance, WideDecimalBinaryExpr, WideDecimalOp, }; use itertools::Itertools; use jni::objects::{Global, JObject}; @@ -325,6 +326,7 @@ pub struct BinaryExprOptions { pub const TEST_EXEC_CONTEXT_ID: i64 = -1; /// The query planner for converting Spark query plans to DataFusion query plans. +#[derive(Clone)] pub struct PhysicalPlanner { // The execution context id of this planner. exec_context_id: i64, @@ -346,6 +348,9 @@ pub struct PhysicalPlanner { /// Task-owned destination for remote shuffle blocks, registered on the driving Spark task /// thread before native planning. Only explicit RSS destinations may use it. shuffle_partition_pusher: Option>, + /// How comparisons treat floating-point operands. `Raw` only while planning a scan's data + /// filters; see [`Self::create_data_filter`]. + float_operands: FloatOperands, } impl Default for PhysicalPlanner { @@ -365,6 +370,7 @@ impl PhysicalPlanner { task_context: None, class_loader: None, shuffle_partition_pusher: None, + float_operands: FloatOperands::Normalize, } } @@ -477,6 +483,11 @@ impl PhysicalPlanner { self } + /// How comparisons treat floating-point operands. + pub fn float_operands(&self) -> FloatOperands { + self.float_operands + } + /// Return session context of this planner. pub fn session_ctx(&self) -> &Arc { &self.session_ctx @@ -700,7 +711,6 @@ impl PhysicalPlanner { SparkCastOptions::new_with_version( eval_mode, &expr.timezone, - expr.allow_incompat, expr.is_spark4_plus, ), spark_expr.expr_id, @@ -793,7 +803,7 @@ impl PhysicalPlanner { "md5" => Ok(Arc::new(Cast::new( func?, DataType::Utf8, - SparkCastOptions::new_without_timezone(EvalMode::Try, true), + SparkCastOptions::new_without_timezone(EvalMode::Try), None, None, ))), @@ -915,8 +925,7 @@ impl PhysicalPlanner { ))) } ExprStruct::ToPrettyString(expr) => { - let mut spark_cast_options = - SparkCastOptions::new(EvalMode::Try, &expr.timezone, true); + let mut spark_cast_options = SparkCastOptions::new(EvalMode::Try, &expr.timezone); let null_string = "NULL"; spark_cast_options.null_string = null_string.to_string(); spark_cast_options.binary_output_style = @@ -1045,6 +1054,8 @@ impl PhysicalPlanner { udf.return_nullable, self.task_context.clone(), self.class_loader.clone(), + self.partition, + self.exec_context_id, ))) } ExprStruct::NativeScalarUdf(call) => { @@ -1127,6 +1138,45 @@ impl PhysicalPlanner { } } + /// Create a data filter that a scan pushes into the Parquet reader, with float operands + /// treated as [`Self::data_filter_float_operands`] decides. + fn create_data_filter( + &self, + spark_expr: &Expr, + input_schema: SchemaRef, + float_operands: FloatOperands, + ) -> Result, ExecutionError> { + let planner = Self { + float_operands, + ..self.clone() + }; + planner.create_expr(spark_expr, input_schema) + } + + /// How a scan's data filters treat float operands. The Parquet reader prunes row groups and + /// pages with them, and with row-level pushdown (`pushdown_filters`) it also drops the rows + /// they reject, which Spark's Filter above the scan then never sees. + /// + /// Without row-level pushdown the filters only prune, and pruning only recognizes a column + /// compared with a literal, so that shape stays a plain comparison of the raw column + /// ([`FloatOperands::Raw`]). With it, every float comparison follows Spark's ordering and + /// gives up pruning: a raw column would drop a stored NaN that Spark matches, such as one with + /// the sign bit set, which Arrow orders below every other value. + fn data_filter_float_operands(&self) -> FloatOperands { + if self + .session_ctx + .copied_config() + .options() + .execution + .parquet + .pushdown_filters + { + FloatOperands::Normalize + } else { + FloatOperands::Raw + } + } + /// Normalize floating-point comparison keys without changing output values: a `FLOAT` or /// `DOUBLE` key, and an array or struct key with a float at any depth, whose order Arrow /// otherwise takes from the raw bits. Sort, Window, and WindowGroupLimit must use identical @@ -1908,10 +1958,17 @@ impl PhysicalPlanner { .cloned() .collect::>(), )); + let float_operands = self.data_filter_float_operands(); common .data_filters .iter() - .map(|expr| self.create_expr(expr, Arc::clone(&filter_schema))) + .map(|expr| { + self.create_data_filter( + expr, + Arc::clone(&filter_schema), + float_operands, + ) + }) .collect() }; @@ -2328,7 +2385,7 @@ impl PhysicalPlanner { let (scans, shuffle_scans, child) = self.create_plan(&children[0], inputs, partition_count)?; - // Create the expression for the array to explode + // Create the expression for the collection to explode let child_expr = if let Some(child_expr) = &explode.child { self.create_expr(child_expr, child.schema())? } else { @@ -2344,6 +2401,22 @@ impl PhysicalPlanner { .name() .to_string(); + // Expose maps as List>, sharing their entries buffers. + // Reuse the list path for outer rows, positions and bounded output batches, + // then flatten the entry struct into Spark's two output columns. + let map_fields = match child_expr.data_type(&child_schema)? { + DataType::Map(entries, _) => match entries.data_type() { + DataType::Struct(fields) => Some(fields.clone()), + _ => unreachable!("Map entries must be a struct"), + }, + _ => None, + }; + let child_expr: Arc = if map_fields.is_some() { + Arc::new(MapEntriesExpr::new(child_expr)) + } else { + child_expr + }; + // Both posexplode variants reference the array twice: once for positions // and once for values. Materialize computed arrays so both references // share one evaluation. A plain Column is already materialized. @@ -2438,11 +2511,22 @@ impl PhysicalPlanner { } }; - output_fields.push(Field::new( - array_field.name(), - element_type, - true, // Element is nullable after unnesting - )); + let struct_unnests = if let Some(fields) = map_fields { + output_fields.extend(fields.iter().map(|field| { + field + .as_ref() + .clone() + .with_nullable(explode.outer || field.is_nullable()) + })); + vec![array_input_index] + } else { + output_fields.push(Field::new( + array_field.name(), + element_type, + true, // Element is nullable after unnesting + )); + vec![] + }; let output_schema = Arc::new(Schema::new(output_fields)); @@ -2472,7 +2556,7 @@ impl PhysicalPlanner { let unnest_exec = Arc::new(ExplodeExec::new( project_exec, list_unnests, - vec![], // No struct columns to unnest + struct_unnests, output_schema, unnest_options, )?); @@ -5824,6 +5908,115 @@ mod tests { assert_eq!(0, filter_exec.additional_native_plans.len()); } + /// Comparisons follow Spark's float ordering without normalizing their operands, except in the + /// data filters that a scan pushes into the Parquet reader, where pruning has to see a plain + /// comparison of the column. + #[test] + fn scan_data_filters_compare_float_columns_directly() { + use datafusion::physical_expr::expressions::BinaryExpr; + use datafusion_comet_spark_expr::SparkComparison; + let double = spark_expression::DataType { + type_id: 6, + type_info: None, + }; + let operand = |expr_struct| Expr { + expr_struct: Some(expr_struct), + query_context: None, + expr_id: None, + }; + let expr = operand(Gt(Box::new(spark_expression::BinaryExpr { + left: Some(Box::new(operand(Bound(spark_expression::BoundReference { + index: 0, + datatype: Some(double.clone()), + })))), + right: Some(Box::new(operand(Literal(spark_expression::Literal { + value: Some(literal::Value::DoubleVal(500.0)), + datatype: Some(double), + is_null: false, + })))), + }))); + let schema = Arc::new(Schema::new(vec![Field::new("d", DataType::Float64, true)])); + let planner = PhysicalPlanner::default(); + let comparison = planner.create_expr(&expr, Arc::clone(&schema)).unwrap(); + let comparison = comparison + .downcast_ref::() + .expect("a Spark comparison"); + assert!(comparison.left().downcast_ref::().is_some()); + let data_filter = planner + .create_data_filter(&expr, schema, planner.data_filter_float_operands()) + .unwrap(); + let data_filter = data_filter + .downcast_ref::() + .expect("a plain comparison"); + assert!(data_filter.left().downcast_ref::().is_some()); + } + + /// With row-level pushdown the Parquet reader drops the rows a data filter rejects, so the + /// filter compares the column in Spark's ordering rather than as it is. A raw column would drop + /// a stored NaN whose bits differ from the normalized literal, and a stored NaN with the sign + /// bit set under any ordering comparison, both of which Spark matches. + #[test] + fn scan_data_filters_follow_spark_ordering_with_row_level_pushdown() { + use arrow::array::{AsArray, BooleanArray}; + use datafusion_comet_spark_expr::{FloatOperands, SparkComparison}; + let double = spark_expression::DataType { + type_id: 6, + type_info: None, + }; + let operand = |expr_struct| Expr { + expr_struct: Some(expr_struct), + query_context: None, + expr_id: None, + }; + let column_and = |value: f64| { + Box::new(spark_expression::BinaryExpr { + left: Some(Box::new(operand(Bound(spark_expression::BoundReference { + index: 0, + datatype: Some(double.clone()), + })))), + right: Some(Box::new(operand(Literal(spark_expression::Literal { + value: Some(literal::Value::DoubleVal(value)), + datatype: Some(double.clone()), + is_null: false, + })))), + }) + }; + // Spark folds `-double('NaN')` into a literal with the sign bit set. + let negative_nan = f64::from_bits(0xfff8_0000_0000_0000); + let schema = Arc::new(Schema::new(vec![Field::new("d", DataType::Float64, true)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Float64Array::from(vec![ + negative_nan, + f64::NAN, + 1.0, + ]))], + ) + .unwrap(); + let config = + SessionConfig::new().set_bool("datafusion.execution.parquet.pushdown_filters", true); + let planner = PhysicalPlanner::new(Arc::new(SessionContext::new_with_config(config)), 0); + assert_eq!( + planner.data_filter_float_operands(), + FloatOperands::Normalize + ); + for (expr, expected) in [ + (operand(Eq(column_and(negative_nan))), [true, true, false]), + (operand(Gt(column_and(0.0))), [true, true, true]), + ] { + let data_filter = planner + .create_data_filter( + &expr, + Arc::clone(&schema), + planner.data_filter_float_operands(), + ) + .unwrap(); + assert!(data_filter.downcast_ref::().is_some()); + let matched = data_filter.evaluate(&batch).unwrap().into_array(3).unwrap(); + assert_eq!(matched.as_boolean(), &BooleanArray::from(expected.to_vec())); + } + } + #[test] fn spark_plan_metrics_hash_join() { let op_scan = create_scan(); @@ -6277,11 +6470,23 @@ mod tests { #[tokio::test] async fn explode_evaluates_array_once_per_batch() { + check_explode_evaluates_collection_once_per_batch(false).await; + } + + #[tokio::test] + async fn explode_evaluates_map_once_per_batch() { + check_explode_evaluates_collection_once_per_batch(true).await; + } + + async fn check_explode_evaluates_collection_once_per_batch(map: bool) { + use arrow::array::{AsArray, MapArray, StructArray}; use arrow::datatypes::Int32Type; use datafusion::common::tree_node::{Transformed, TreeNode}; use datafusion::logical_expr::{create_udf, Volatility}; use datafusion::physical_plan::projection::ProjectionExec; - use spark_expression::data_type::{data_type_info::DatatypeStruct, DataTypeInfo, ListInfo}; + use spark_expression::data_type::{ + data_type_info::DatatypeStruct, DataTypeInfo, ListInfo, MapInfo, + }; let array_type = spark_expression::DataType { type_id: 14, @@ -6300,6 +6505,55 @@ mod tests { Some(vec![Some(20)]), ])) as ArrayRef; + let (array_type, arrays) = if map { + // Slice away a leading entry to exercise non-zero map offsets. + let values = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(99)]), + Some(vec![Some(10), None]), + Some(vec![]), + None, + Some(vec![Some(20)]), + ]) + .slice(1, 4); + let fields: Fields = vec![ + Field::new("key", DataType::Int32, false) + .with_metadata([("PARQUET:field_id".to_string(), "10".to_string())].into()), + Field::new("value", DataType::Int32, true) + .with_metadata([("PARQUET:field_id".to_string(), "11".to_string())].into()), + ] + .into(); + let entries = StructArray::new( + fields.clone(), + vec![ + Arc::new(Int32Array::from(vec![99, 1, 2, 3])), + Arc::clone(values.values()), + ], + None, + ); + let maps = MapArray::new( + Arc::new(Field::new("entries", DataType::Struct(fields), false)), + values.offsets().clone(), + entries, + values.nulls().cloned(), + false, + ); + let map_type = spark_expression::DataType { + type_id: 15, + type_info: Some(Box::new(DataTypeInfo { + datatype_struct: Some(DatatypeStruct::Map(Box::new(MapInfo { + key_type: Some(Box::new(create_proto_datatype())), + value_type: Some(Box::new(create_proto_datatype())), + value_contains_null: true, + key_field_id: Some(10), + value_field_id: Some(11), + }))), + })), + }; + (map_type, Arc::new(maps) as ArrayRef) + } else { + (array_type, arrays) + }; + for outer in [false, true] { for position in [false, true] { for computed in [false, true] { @@ -6378,20 +6632,21 @@ mod tests { let results = collect(native_plan.execute(0, task_ctx).unwrap()) .await .unwrap(); - let context = - format!("outer={outer}, position={position}, computed={computed}"); + let context = format!( + "outer={outer}, position={position}, computed={computed}, map={map}" + ); assert_eq!( calls.load(Ordering::Relaxed), if computed { 2 } else { 0 }, "{context}" ); // The array is pre-projected only to share one evaluation between the - // `pos` and `value` references, so only a computed child needs it. `outer` + // `pos` and `value` references, so a computed child or map wrapper needs it. `outer` // does not, since it is now a `UnnestOptions` mode rather than a wrapper // expression around the child. assert_eq!( projections, - 1 + usize::from(position && computed), + 1 + usize::from(position && (computed || map)), "{context}" ); let expected_values = if outer { @@ -6411,6 +6666,23 @@ mod tests { }) .collect(); assert_eq!(values, expected_values.repeat(2), "{context}"); + if map { + let expected_keys = if outer { + vec![Some(1), Some(2), None, None, Some(3)] + } else { + vec![Some(1), Some(2), Some(3)] + }; + let keys: Vec<_> = results + .iter() + .flat_map(|batch| { + batch + .column(batch.num_columns() - 2) + .as_primitive::() + .iter() + }) + .collect(); + assert_eq!(keys, expected_keys.repeat(2), "{context}"); + } if position { let expected_positions = if outer { vec![Some(0), Some(1), None, None, Some(0)] @@ -6740,7 +7012,7 @@ mod tests { .with_table_parquet_options(TableParquetOptions::new()), ) as Arc; - let spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); let expr_adapter_factory: Arc = Arc::new( SparkPhysicalExprAdapterFactory::new(spark_parquet_options, None), diff --git a/native/core/src/execution/planner/macros.rs b/native/core/src/execution/planner/macros.rs index d146f6633bb..5c11770f09d 100644 --- a/native/core/src/execution/planner/macros.rs +++ b/native/core/src/execution/planner/macros.rs @@ -98,6 +98,7 @@ macro_rules! binary_expr_builder { $operator, right, input_schema.as_ref(), + planner.float_operands(), ) .map_err(Into::into) } diff --git a/native/core/src/execution/planner/write.rs b/native/core/src/execution/planner/write.rs index 1a3f9674cd1..5665939afb2 100644 --- a/native/core/src/execution/planner/write.rs +++ b/native/core/src/execution/planner/write.rs @@ -107,9 +107,6 @@ impl OperatorBuilder for ParquetWriterBuilder { let parquet_writer = Arc::new(ParquetWriterExec::try_new( Arc::clone(&child.native_plan), writer.output_path.clone(), - writer.work_dir.clone(), - writer.job_id.clone(), - writer.task_attempt_id, codec, planner.partition(), writer.column_names.clone(), diff --git a/native/core/src/execution/spark_config.rs b/native/core/src/execution/spark_config.rs index d29bf6f6836..1de88ccbeec 100644 --- a/native/core/src/execution/spark_config.rs +++ b/native/core/src/execution/spark_config.rs @@ -26,6 +26,8 @@ pub(crate) const COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED: &str = "spark.comet.parquet.rowFilterPushdown.enabled"; pub(crate) const COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED: &str = "spark.comet.exec.aggregate.skipPartial.enabled"; +pub(crate) const COMET_EXEC_SPILL_COMPRESSION_CODEC: &str = + "spark.comet.exec.spill.compression.codec"; pub(crate) const SPARK_EXECUTOR_CORES: &str = "spark.executor.cores"; /// Comet configs read through this trait must be resolved by the JVM first: diff --git a/native/core/src/lib.rs b/native/core/src/lib.rs index b712a65c115..5eb56aaec01 100644 --- a/native/core/src/lib.rs +++ b/native/core/src/lib.rs @@ -31,6 +31,7 @@ extern crate datafusion_comet_jni_bridge; use jni::{ objects::{JClass, JString}, + sys::{jboolean, jstring, JNI_FALSE}, EnvUnowned, }; use log::info; @@ -52,6 +53,9 @@ pub mod jvm_bridge { use errors::{try_unwrap_or_throw, CometError, CometResult}; +use crate::cloud::s3::credential_bridge::AccessMode; +use crate::execution::operators::iceberg_common::builtin_storage_schemes; + pub mod alloc_accounting; pub mod cloud; pub mod comet_native_udf_bridge; @@ -270,6 +274,27 @@ pub fn tzdata_version() -> &'static str { chrono_tz::IANA_TZDB_VERSION } +/// JNI: the comma-joined built-in storage schemes the native Iceberg scan (`for_write` false) or +/// write (`for_write` true) path admits. This is the source of truth the JVM read and write gates +/// load. Opt-in S3-compliant alias schemes are not answered here; the JVM adds them from catalog +/// properties. +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_NativeBase_icebergStorageSchemes( + env: EnvUnowned, + _: JClass, + for_write: jboolean, +) -> jstring { + try_unwrap_or_throw(&env, |env| { + let access_mode = if for_write != JNI_FALSE { + AccessMode::Write + } else { + AccessMode::Read + }; + let joined = builtin_storage_schemes(access_mode).join(","); + Ok(env.new_string(joined)?.into_raw()) + }) +} + // Creates a default log4rs config, which logs to console with log level. fn default_logger_config(log_level: &str) -> CometResult { let console_append = ConsoleAppender::builder() diff --git a/native/core/src/parquet/cast_column.rs b/native/core/src/parquet/cast_column.rs index ad7b909768b..3a76e89916f 100644 --- a/native/core/src/parquet/cast_column.rs +++ b/native/core/src/parquet/cast_column.rs @@ -393,7 +393,7 @@ mod tests { let schema = Schema::new(vec![Arc::clone(&input_field)]); let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(struct_arr)]).unwrap(); - let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); opts.use_field_id = true; let col_expr: Arc = Arc::new(Column::new("s", 0)); @@ -466,7 +466,7 @@ mod tests { let expr: Arc = Arc::new(Column::new("ts", 0)); let cast_expr = CometCastColumnExpr::try_new(expr, input_field, target_field, None) .unwrap() - .with_parquet_options(SparkParquetOptions::new(eval_mode, "UTC", false)); + .with_parquet_options(SparkParquetOptions::new(eval_mode, "UTC")); let input = TimestampMillisecondArray::from(vec![Some(1_234), Some(-1_234), None]) .with_timezone_opt(source_tz.clone()); diff --git a/native/core/src/parquet/eager_page_index_reader_factory.rs b/native/core/src/parquet/eager_page_index_reader_factory.rs index dbc7d0f3b80..dbde337306e 100644 --- a/native/core/src/parquet/eager_page_index_reader_factory.rs +++ b/native/core/src/parquet/eager_page_index_reader_factory.rs @@ -52,9 +52,8 @@ //! `ignoreMissing` is set, and it walks the raw `MessageType` to decide. That walk has to run //! over the Parquet schema rather than the Arrow schema the schema adapter is handed later, //! because an id on a repeated `list` or `key_value` group, or on the message root, never -//! reaches an Arrow field. Until apache/datafusion#24790, which is not in DataFusion 55.1.0, -//! the INT96 coercion also rebuilt container fields without their metadata, so a struct id could -//! vanish on the way to Arrow as well. +//! reaches an Arrow field. DataFusion 55.2.0 preserves container metadata through INT96 coercion +//! (apache/datafusion#24790), so struct ids no longer vanish on the way to Arrow. //! //! The second is the Variant footer rewrite, `with_spark_arrow_schema`, which replaces the //! Arrow schema hint in the footer for scans that project Variant. @@ -1341,7 +1340,7 @@ mod tests { false, )])); let mut options = - SparkParquetOptions::new(datafusion_comet_spark_expr::EvalMode::Legacy, "UTC", false); + SparkParquetOptions::new(datafusion_comet_spark_expr::EvalMode::Legacy, "UTC"); options.case_sensitive = false; let factory = EagerPageIndexReaderFactory::new( store, diff --git a/native/core/src/parquet/objectstore/s3_blob_fs_support.rs b/native/core/src/parquet/objectstore/s3_blob_fs_support.rs index f455f8c3849..ab09e7d24c3 100644 --- a/native/core/src/parquet/objectstore/s3_blob_fs_support.rs +++ b/native/core/src/parquet/objectstore/s3_blob_fs_support.rs @@ -98,8 +98,10 @@ pub(crate) fn normalize_object_store_url( /// True if `scheme` is a configured s3-compliant alias (`fs.comet.s3Compliant.schemes`; empty or /// unset means none). These route to `AmazonS3` but are not recognized by /// `ObjectStoreScheme::parse`. `s3a` is excluded -- object_store knows it and callers special-case -/// it. Shared by the Parquet normalizer above and the Iceberg storage-factory gate -/// (`iceberg_common::storage_factory_for`) so both admit the same schemes. +/// it. Case-insensitive, which is safe for the Parquet normalizer above because it rewrites the +/// alias to `s3://` before anything opens the path. The Iceberg storage-factory gate opens the +/// recorded location as written, so it wraps this in `iceberg_common::is_iceberg_alias_scheme`, +/// which additionally requires the scheme to be lowercase. pub(crate) fn is_s3_compliant_alias_scheme( scheme: &str, object_store_configs: &HashMap, diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index 8149c3656b7..fba2d6fc13e 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -351,8 +351,7 @@ fn get_options( // distinguish INT96-derived TimestampLTZ from a true TimestampNTZ source // and apply the pre-Spark-4 SPARK-36182 rejection (#4219). table_parquet_options.global.coerce_int96_tz = Some("UTC".to_string()); - let mut spark_parquet_options = - SparkParquetOptions::new(EvalMode::Legacy, session_timezone, false); + let mut spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, session_timezone); spark_parquet_options.allow_cast_unsigned_ints = true; spark_parquet_options.case_sensitive = case_sensitive; spark_parquet_options.return_null_struct_if_all_fields_missing = @@ -417,6 +416,8 @@ mod tests { true, None, None, + 0, + 0, )); let literal: Arc = Arc::new(Literal::new(ScalarValue::Int32(Some(1)))); let nested: Arc = @@ -1346,6 +1347,79 @@ mod tests { assert!(scan_metric(&scan, "scan_io_metadata_bytes") >= bloom_bytes); } + /// A bloom filter holds the bits of each value, so `-0.0` and `0.0` are separate entries in it, + /// while Spark's `=` matches either zero. A data filter that only prunes (no row-level + /// pushdown) probes the filter for both zeros, so the row group holding `-0.0` and no `0.0` + /// is read for `d = 0.0` and for `d = -0.0`, and a value the filter does not hold still + /// prunes it. The column statistics cover all three values, so only the bloom filter prunes. + #[tokio::test] + async fn bloom_filter_pruning_matches_either_signed_zero() { + use arrow::array::Float64Array; + use datafusion::physical_plan::metrics::MetricValue; + use datafusion_comet_spark_expr::{spark_comparison, FloatOperands}; + + let schema = Arc::new(Schema::new(vec![Field::new("d", DataType::Float64, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Float64Array::from(vec![-1.0, -0.0, 1.0]))], + ) + .unwrap(); + let filename = get_temp_filename() + .as_path() + .as_os_str() + .to_str() + .unwrap() + .to_string(); + let props = WriterProperties::builder() + .set_bloom_filter_enabled(true) + .build(); + let file = File::create(&filename).unwrap(); + let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let session_ctx = Arc::new(SessionContext::new()); + for (value, read) in [(-0.0, true), (0.0, true), (0.5, false)] { + let filter = spark_comparison( + Arc::new(Column::new("d", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Float64(Some(value)))), + &schema, + FloatOperands::Raw, + ) + .unwrap(); + let scan = init_test_scan( + Arc::clone(&schema), + Arc::clone(&schema), + PartitionedFile::from_path(filename.clone()).unwrap(), + None, + Some(vec![filter]), + &session_ctx, + ); + let mut stream = scan.execute(0, session_ctx.task_ctx()).unwrap(); + let mut rows = 0; + while let Some(batch) = stream.next().await { + rows += batch.unwrap().num_rows(); + } + let bloom_filter = scan + .metrics() + .unwrap() + .sum_by_name("row_groups_pruned_bloom_filter"); + let Some(MetricValue::PruningMetrics { + pruning_metrics, .. + }) = bloom_filter + else { + panic!("missing bloom filter pruning metrics: {bloom_filter:?}"); + }; + let expected = if read { (0, 1, 3) } else { (1, 0, 0) }; + assert_eq!( + (pruning_metrics.pruned(), pruning_metrics.matched(), rows), + expected, + "d = {value:?}" + ); + } + } + #[tokio::test] async fn reports_projected_and_predicate_pruned_data_io() { let (filename, schema) = write_scan_io_fixture(); diff --git a/native/core/src/parquet/parquet_support.rs b/native/core/src/parquet/parquet_support.rs index e8987a79da6..cc17429a72b 100644 --- a/native/core/src/parquet/parquet_support.rs +++ b/native/core/src/parquet/parquet_support.rs @@ -87,8 +87,6 @@ pub struct SparkParquetOptions { /// session local timezone by an analyzer in Spark. // TODO we should change timezone to Tz to avoid repeated parsing pub timezone: String, - /// Allow casts that are supported but not guaranteed to be 100% compatible - pub allow_incompat: bool, /// Support casting unsigned ints to signed ints (used by Parquet SchemaAdapter) pub allow_cast_unsigned_ints: bool, /// Whether to read dates/timestamps that were written in the legacy hybrid Julian + Gregorian calendar as it is. If false, throw exceptions instead. If the spark type is TimestampNTZ, this should be true. @@ -121,11 +119,10 @@ pub struct SparkParquetOptions { } impl SparkParquetOptions { - pub fn new(eval_mode: EvalMode, timezone: &str, allow_incompat: bool) -> Self { + pub fn new(eval_mode: EvalMode, timezone: &str) -> Self { Self { eval_mode, timezone: timezone.to_string(), - allow_incompat, allow_cast_unsigned_ints: false, use_legacy_date_timestamp_or_ntz: false, case_sensitive: false, @@ -137,11 +134,10 @@ impl SparkParquetOptions { } } - pub fn new_without_timezone(eval_mode: EvalMode, allow_incompat: bool) -> Self { + pub fn new_without_timezone(eval_mode: EvalMode) -> Self { Self { eval_mode, timezone: "".to_string(), - allow_incompat, allow_cast_unsigned_ints: false, use_legacy_date_timestamp_or_ntz: false, case_sensitive: false, @@ -1405,7 +1401,7 @@ mod tests { let err = spark_parquet_convert( ColumnarValue::Array(Arc::new(array)), &to_type, - &SparkParquetOptions::new(EvalMode::Legacy, "UTC", false), + &SparkParquetOptions::new(EvalMode::Legacy, "UTC"), ) .expect_err("array -> int must be an error"); assert!( @@ -1513,7 +1509,7 @@ mod tests { use datafusion_comet_spark_expr::EvalMode; use std::sync::Arc; - let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); let overflow_millis = 9_223_372_036_854_776_i64; let millis: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![ Some(overflow_millis), @@ -1633,7 +1629,7 @@ mod tests { )); let target = DataType::Struct(target_fields.into()); let input: ArrayRef = Arc::new(input); - let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); let output = parquet_convert_array(Arc::clone(&input), &target, &options).unwrap(); assert_eq!(output.data_type(), &target); assert!(output.is_null(0)); @@ -1682,7 +1678,7 @@ mod tests { for timezone in [None::>, Some(Arc::from("UTC"))] { for overflow in [i64::MAX, i64::MIN] { - let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); let millis: ArrayRef = Arc::new( TimestampMillisecondArray::from(vec![overflow, 7, overflow]) .with_timezone_opt(timezone.clone()), @@ -1822,7 +1818,7 @@ mod tests { use datafusion_comet_spark_expr::EvalMode; use std::sync::Arc; - let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); for timezone in [None::>, Some(Arc::from("UTC"))] { for overflow in [i64::MIN, i64::MAX] { let values: ArrayRef = Arc::new( @@ -1918,7 +1914,7 @@ mod tests { use datafusion_comet_spark_expr::EvalMode; use std::sync::Arc; - let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); let values: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![ i64::MAX, 7, @@ -2125,7 +2121,7 @@ mod tests { )); let to_type = struct_type_with_field_id(vec![("f", DataType::Int32, 1)]); - let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); opts.use_field_id = true; let err = parquet_convert_array(from, &to_type, &opts).unwrap_err(); @@ -2165,7 +2161,7 @@ mod tests { )); let to_type = struct_type_with_field_id(vec![("f", DataType::Int32, 2)]); - let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); opts.use_field_id = true; let result = parquet_convert_array(from, &to_type, &opts).unwrap(); diff --git a/native/core/src/parquet/schema_adapter.rs b/native/core/src/parquet/schema_adapter.rs index 6bcc8d3518b..495eeb545ec 100644 --- a/native/core/src/parquet/schema_adapter.rs +++ b/native/core/src/parquet/schema_adapter.rs @@ -1475,7 +1475,6 @@ impl SparkPhysicalExprAdapter { let mut cast_options = SparkCastOptions::new( self.parquet_options.eval_mode, &self.parquet_options.timezone, - self.parquet_options.allow_incompat, ); cast_options.allow_cast_unsigned_ints = self.parquet_options.allow_cast_unsigned_ints; cast_options.is_adapting_schema = true; @@ -2202,7 +2201,7 @@ pub(crate) mod test { let required_schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int64, false)])); - let mut spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); spark_parquet_options.allow_type_promotion = false; let expr_adapter_factory: Arc = Arc::new( @@ -2247,7 +2246,7 @@ pub(crate) mod test { let required_schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int64, false)])); - let mut spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); spark_parquet_options.allow_type_promotion = false; let expr_adapter_factory: Arc = Arc::new( @@ -2375,7 +2374,7 @@ pub(crate) mod test { batch: &RecordBatch, required_schema: SchemaRef, ) -> Result { - let mut spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); spark_parquet_options.allow_cast_unsigned_ints = true; let mut stream = scan_parquet(batch, required_schema, spark_parquet_options)?; stream.next().await.unwrap() @@ -2890,7 +2889,7 @@ pub(crate) mod test { #[test] fn nested_dictionary_containers_are_checked_by_value_type() -> Result<(), DataFusionError> { - let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); let physical_struct = DataType::Struct(Fields::from(vec![ Field::new("x", DataType::Int64, true), Field::new("unused", DataType::Int32, true), @@ -2925,7 +2924,7 @@ pub(crate) mod test { #[test] fn nested_map_shape_mismatch_is_rejected() -> Result<(), DataFusionError> { - let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); let valid = map_type(DataType::Int64); let DataType::Map(entries, _) = &valid else { unreachable!() @@ -3048,7 +3047,7 @@ pub(crate) mod test { let required_schema = struct_schema(vec![ Field::new("b", DataType::Int32, true).with_metadata(id_meta("1")) ]); - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.use_field_id = true; let mut stream = scan_parquet(&batch, required_schema, options)?; let err = stream @@ -3075,7 +3074,7 @@ pub(crate) mod test { Arc::new(Int32Array::from(vec![1, 2, 3])), )?; let required_schema = struct_schema(vec![Field::new("x", DataType::Int64, true)]); - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.allow_type_promotion = false; let mut stream = scan_parquet(&batch, required_schema, options)?; let err = stream @@ -3101,7 +3100,7 @@ pub(crate) mod test { Arc::new(Int32Array::from(Vec::::new())), )?; let required_schema = struct_schema(vec![Field::new("x", DataType::Int64, true)]); - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.allow_type_promotion = false; let mut stream = scan_parquet(&batch, required_schema, options)?; while let Some(batch) = stream.next().await { @@ -3486,7 +3485,7 @@ pub(crate) mod test { Arc::new(Int32Array::from(vec![1, 2, 3])), )?; let required_schema = struct_schema(vec![Field::new("x", DataType::Int64, true)]); - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.allow_type_promotion = true; let mut stream = scan_parquet(&batch, required_schema, options)?; let result = stream.next().await.unwrap()?; @@ -3579,7 +3578,7 @@ pub(crate) mod test { Field::new("a", DataType::Int64, nullable), Field::new("b", DataType::Int64, false), ])); - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.case_sensitive = true; let adapter = SparkPhysicalExprAdapterFactory::new(options, None) .create(logical, physical) @@ -3613,7 +3612,7 @@ pub(crate) mod test { Arc::new(Field::new("dup", DataType::Int64, true)), Arc::new(Field::new("dup", DataType::Int64, true)), ]; - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.case_sensitive = true; let error = super::match_struct_fields(&fields, &fields[..1], &options) .expect_err("selected duplicate child must fail") @@ -3648,7 +3647,7 @@ pub(crate) mod test { default.rewrite(Arc::clone(&column)).is_err(), "fixture must reach the default-adapter fallback" ); - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.case_sensitive = true; let adapter = SparkPhysicalExprAdapterFactory::new(options, None) .create(Arc::clone(&logical), Arc::clone(&physical)) @@ -3714,7 +3713,7 @@ pub(crate) mod test { if mode == "missing" { requested.push(Field::new("missing", DataType::Int64, true)); } - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.case_sensitive = true; let mut stream = scan_parquet(&batch, struct_schema(requested), options).unwrap(); let error = stream @@ -3816,7 +3815,7 @@ pub(crate) mod test { reader.schema().field(0).data_type(), required.field(0).data_type() ); - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.case_sensitive = true; let mut stream = scan_parquet(&batch, required, options).unwrap(); let error = stream @@ -3859,7 +3858,7 @@ pub(crate) mod test { // Read with case-insensitive mode, requesting column "b" which matches both "B" and "b" let required_schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Int32, false)])); - let mut spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); spark_parquet_options.case_sensitive = false; let expr_adapter_factory: Arc = Arc::new( @@ -3919,7 +3918,7 @@ pub(crate) mod test { None, )?)); let defaults = HashMap::from([(Column::new("missing", 0), default)]); - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.case_sensitive = false; let adapter = SparkPhysicalExprAdapterFactory::new(options, Some(defaults)) .create(logical, physical)?; @@ -4014,7 +4013,7 @@ pub(crate) mod test { Field::new("ω", DataType::Int32, true).with_metadata(id_meta("2")), ])); - let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); opts.case_sensitive = false; opts.use_field_id = true; let adapter = SparkPhysicalExprAdapterFactory::new(opts, None) @@ -4045,7 +4044,7 @@ pub(crate) mod test { let physical = Arc::new(Schema::new(vec![ Field::new("MÜNCHEN", storage, true).with_extension_type(VariantType) ])); - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.case_sensitive = false; let adapter = SparkPhysicalExprAdapterFactory::new(options, None) .create(Arc::clone(&logical), Arc::clone(&physical)) @@ -4104,7 +4103,7 @@ pub(crate) mod test { .is_some()); } let adapter = SparkPhysicalExprAdapterFactory::new( - SparkParquetOptions::new(EvalMode::Legacy, "UTC", false), + SparkParquetOptions::new(EvalMode::Legacy, "UTC"), None, ) .create(logical, physical) @@ -4140,7 +4139,7 @@ pub(crate) mod test { Field::new("other", storage.clone(), true).with_metadata(id_meta("1")), Field::new("__comet_unmatched_field_id_1", DataType::Binary, true), ])); - let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); options.case_sensitive = case_sensitive; options.use_field_id = true; let adapter = SparkPhysicalExprAdapterFactory::new(options, None) @@ -4210,7 +4209,7 @@ pub(crate) mod test { )); let target_type = DataType::List(Arc::clone(&to_item_field)); - let spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); let comet_result = spark_parquet_convert( ColumnarValue::Array(Arc::clone(&list_array)), &target_type, @@ -4321,7 +4320,7 @@ pub(crate) mod test { } fn default_options() -> SparkParquetOptions { - SparkParquetOptions::new(EvalMode::Legacy, "UTC", false) + SparkParquetOptions::new(EvalMode::Legacy, "UTC") } /// Dropping a struct field by exact name, including through nested struct-in-struct and @@ -4683,7 +4682,7 @@ pub(crate) mod test { )) as Arc; let batch = RecordBatch::try_new(file_schema, vec![col]).unwrap(); - let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false); + let mut opts = SparkParquetOptions::new(EvalMode::Legacy, "UTC"); opts.use_field_id = true; let err = match scan_parquet(&batch, required_schema, opts) { diff --git a/native/jni-bridge/src/comet_udf_bridge.rs b/native/jni-bridge/src/comet_udf_bridge.rs index bbd7465d007..6532fe04035 100644 --- a/native/jni-bridge/src/comet_udf_bridge.rs +++ b/native/jni-bridge/src/comet_udf_bridge.rs @@ -42,7 +42,7 @@ impl<'a> CometUdfBridge<'a> { JNIString::new(Self::JVM_CLASS), jni::jni_str!("evaluate"), jni::jni_sig!( - "(Ljava/lang/String;[J[JJJILorg/apache/spark/TaskContext;Ljava/lang/ClassLoader;)V" + "(Ljava/lang/String;[J[JJJIIJLorg/apache/spark/TaskContext;Ljava/lang/ClassLoader;)V" ), )?, method_evaluate_ret: ReturnType::Primitive(Primitive::Void), diff --git a/native/operators/benches/explode.rs b/native/operators/benches/explode.rs index ef14e722548..c5e624fa641 100644 --- a/native/operators/benches/explode.rs +++ b/native/operators/benches/explode.rs @@ -28,11 +28,13 @@ use std::sync::Arc; -use arrow::array::{Array, ArrayRef, Int32Array, Int64Array, ListArray, StringArray, StructArray}; +use arrow::array::{ + Array, ArrayRef, AsArray, Int32Array, Int64Array, ListArray, StringArray, StructArray, +}; use arrow::buffer::{NullBuffer, OffsetBuffer}; -use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; +use arrow::datatypes::{DataType, Field, Fields, Int64Type, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; -use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; use datafusion::common::{NullHandling, UnnestOptions}; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::execution::TaskContext; @@ -264,6 +266,173 @@ fn run(runtime: &Runtime, plan: &Arc, ctx: &Arc) assert!(!batches.is_empty()); } +// The inner generator in CometExplodeBenchmark's chained map case carries k, region and +// platform through five entries per map value. The existing carried-column group only has +// Int64 columns, so it cannot expose the string gathers in this shape. +const CHAINED_FAN_OUT: usize = 5; +const K: usize = 0; +const REGION: usize = 1; +const PLATFORM: usize = 2; + +fn chained_entries(start: usize) -> ListArray { + let entry_fields = Fields::from(vec![ + Field::new("type", DataType::Utf8, true), + Field::new("ts", DataType::Int64, true), + Field::new("page", DataType::Utf8, true), + Field::new("source", DataType::Utf8, true), + ]); + let entry_start = start * CHAINED_FAN_OUT; + let entry_count = ROWS_PER_BATCH * CHAINED_FAN_OUT; + let entries = StructArray::new( + entry_fields, + vec![ + Arc::new(StringArray::from_iter_values( + (0..entry_count).map(|i| format!("t_{}", i % CHAINED_FAN_OUT)), + )), + Arc::new(Int64Array::from_iter_values( + (entry_start..entry_start + entry_count).map(|i| i as i64), + )), + Arc::new(StringArray::from_iter_values( + (entry_start..entry_start + entry_count).map(|i| format!("page_{i}")), + )), + Arc::new(StringArray::from_iter_values( + (0..entry_count).map(|i| format!("src_{}", i % 4)), + )), + ], + None, + ); + ListArray::new( + Arc::new(Field::new("item", entries.data_type().clone(), true)), + OffsetBuffer::from_lengths(std::iter::repeat_n(CHAINED_FAN_OUT, ROWS_PER_BATCH)), + Arc::new(entries), + None, + ) +} + +fn chained_plan(carried: &[usize], nullable: bool) -> (Arc, Vec) { + let batches: Vec<_> = (0..BATCHES) + .map(|batch| { + let start = batch * ROWS_PER_BATCH; + let list = chained_entries(start); + let k = Int64Array::from_iter( + (start..start + ROWS_PER_BATCH) + .map(|i| (!nullable || !i.is_multiple_of(10)).then_some((i / 4) as i64)), + ); + let region = + StringArray::from_iter((start..start + ROWS_PER_BATCH).map(|i| { + (!nullable || !i.is_multiple_of(10)).then(|| format!("r_{}", i / 4 % 8)) + })); + let platform = StringArray::from_iter_values( + (start..start + ROWS_PER_BATCH).map(|i| format!("p_{}", i % 4 + 1)), + ); + let payload: [ArrayRef; 3] = [Arc::new(k), Arc::new(region), Arc::new(platform)]; + let names = ["k", "region", "platform"]; + let mut fields = vec![Field::new("entries", list.data_type().clone(), true)]; + let mut columns: Vec = vec![Arc::new(list)]; + for &i in carried { + fields.push(Field::new( + names[i], + payload[i].data_type().clone(), + i != PLATFORM, + )); + columns.push(Arc::clone(&payload[i])); + } + RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap() + }) + .collect(); + let schema = batches[0].schema(); + let output_fields = std::iter::once(Arc::new(Field::new( + "entry", + batches[0].column(0).as_list::().value_type(), + true, + ))) + .chain(schema.fields().iter().skip(1).cloned()) + .collect::>(); + let source = + MemorySourceConfig::try_new_exec(std::slice::from_ref(&batches), schema, None).unwrap(); + let plan = ExplodeExec::new( + source, + vec![ListUnnest { + index_in_input_schema: 0, + depth: 1, + }], + vec![], + Arc::new(Schema::new(output_fields)), + UnnestOptions::new().with_null_handling(NullHandling::Drop), + ) + .unwrap(); + (Arc::new(plan), batches) +} + +fn verify_chained( + runtime: &Runtime, + plan: &Arc, + ctx: &Arc, + input: &[RecordBatch], + carried: &[usize], + nullable: bool, +) { + let output = runtime + .block_on(collect(plan.execute(0, Arc::clone(ctx)).unwrap())) + .unwrap(); + let mut row = 0; + let mut null_counts = vec![0; carried.len()]; + for batch in output { + assert!(batch.num_rows() <= ROWS_PER_BATCH); + let entries = batch.column(0).as_struct(); + let source_batch = row / (ROWS_PER_BATCH * CHAINED_FAN_OUT); + let source_list = input[source_batch].column(0).as_list::(); + let source_entries = source_list.values().as_struct(); + // String slicing retains the full values buffer. A gather would allocate a new one. + let page = entries.column(2).as_string::(); + let source_page = source_entries.column(2).as_string::(); + assert_eq!( + page.value_data().as_ptr(), + source_page.value_data().as_ptr() + ); + let ts = entries.column(1).as_primitive::(); + for (position, &column) in carried.iter().enumerate() { + let array = batch.column(position + 1); + null_counts[position] += array.null_count(); + for j in 0..batch.num_rows() { + let input_row = (row + j) / CHAINED_FAN_OUT; + let is_null = nullable && column != PLATFORM && input_row.is_multiple_of(10); + assert_eq!(array.is_null(j), is_null); + if is_null { + continue; + } + if column == K { + assert_eq!( + array.as_primitive::().value(j), + (input_row / 4) as i64 + ); + } else { + let expected = if column == REGION { + format!("r_{}", input_row / 4 % 8) + } else { + format!("p_{}", input_row % 4 + 1) + }; + let value = array.as_string::().value(j); + assert_eq!(value, expected); + } + } + } + for j in 0..batch.num_rows() { + assert_eq!(ts.value(j), (row + j) as i64); + } + row += batch.num_rows(); + } + assert_eq!(row, ROWS_PER_BATCH * BATCHES * CHAINED_FAN_OUT); + for (position, &column) in carried.iter().enumerate() { + let expected = if nullable && column != PLATFORM { + (ROWS_PER_BATCH * BATCHES).div_ceil(10) * CHAINED_FAN_OUT + } else { + 0 + }; + assert_eq!(null_counts[position], expected); + } +} + fn criterion_benchmark(c: &mut Criterion) { let runtime = Runtime::new().unwrap(); let ctx = Arc::new( @@ -320,6 +489,24 @@ fn criterion_benchmark(c: &mut Criterion) { } } group.finish(); + + let mut group = c.benchmark_group("explode_chained_carry"); + group.throughput(Throughput::Elements( + (ROWS_PER_BATCH * BATCHES * CHAINED_FAN_OUT) as u64, + )); + for (name, carried, nullable) in [ + ("none", &[][..], false), + ("k", &[K][..], false), + ("region", &[REGION][..], false), + ("platform", &[PLATFORM][..], false), + ("full", &[K, REGION, PLATFORM][..], false), + ("full_nulls", &[K, REGION, PLATFORM][..], true), + ] { + let (plan, input) = chained_plan(carried, nullable); + verify_chained(&runtime, &plan, &ctx, &input, carried, nullable); + group.bench_function(name, |b| b.iter(|| run(&runtime, &plan, &ctx))); + } + group.finish(); } criterion_group!(benches, criterion_benchmark); diff --git a/native/operators/src/explode.rs b/native/operators/src/explode.rs index 5af1e052d7b..cc2159db142 100644 --- a/native/operators/src/explode.rs +++ b/native/operators/src/explode.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! A fork of DataFusion's `UnnestExec`, kept for two unnesting kernels Comet has specialized. +//! A fork of DataFusion's `UnnestExec`, kept for Comet's specialized unnesting kernels. //! //! # Why this fork exists //! @@ -23,16 +23,20 @@ //! however many rows the unnesting produced, and never consulted //! `datafusion.execution.batch_size`. That fix is now upstream //! (apache/datafusion#24384, in DataFusion 55.1.0), so it is no longer the reason to keep -//! the fork. What is left are two performance paths that have not been upstreamed: +//! the fork. What is left are performance improvements that have not been upstreamed: //! //! * `list_output_lens`, which computes the per-row output lengths of a single `List` column //! in one pass over the offsets instead of chaining six arrow kernels; //! * the contiguous-run fast path in `unnest_list_array`, which returns a slice of the child //! values instead of gathering them, and the buffer fills in `create_take_indices`. +//! * retaining input arrays as unnest placeholders and skipping take indices when no column +//! consumes them, avoiding allocations whose contents never reach the output. +//! * repeating carried short strings by input row, with fixed-size byte copies instead of +//! a variable-size copy for every output row. //! //! # Deleting this file //! -//! Upstream those two paths, then delete this module and go back to +//! Upstream these improvements, then delete this module and go back to //! `datafusion::physical_plan::unnest::UnnestExec` in the planner. Deleting it before that //! would regress `explode`, so measure with `native/operators/benches/explode.rs` first. //! @@ -51,15 +55,16 @@ //! * dropping upstream's `ListUnnest` declaration in favor of importing the public one; //! * `find_longest_length` applies the empty-list bump once after the row-wise maximum rather //! than once per array, which is equivalent because `max` is associative; -//! * `list_output_lens` and the contiguous-run fast path described above. +//! * the performance improvements described above. //! //! `ExplodeExec` and `ExplodeStream` were always Comet's own. use arrow::array::{ - new_null_array, Array, ArrayRef, AsArray, BooleanBufferBuilder, FixedSizeListArray, Int64Array, - LargeListArray, LargeListViewArray, ListArray, ListViewArray, PrimitiveArray, Scalar, - StructArray, + Array, ArrayRef, AsArray, BooleanBufferBuilder, FixedSizeListArray, Int64Array, LargeListArray, + LargeListViewArray, ListArray, ListViewArray, NullBufferBuilder, PrimitiveArray, Scalar, + StringArray, StructArray, }; +use arrow::buffer::OffsetBuffer; use arrow::compute::kernels::length::length; use arrow::compute::kernels::zip::zip; use arrow::compute::{cast, is_not_null, kernels, sum}; @@ -685,8 +690,6 @@ fn list_unnest_at_level( let unnested_temp_arrays = unnest_list_arrays(arrs_to_unnest.as_ref(), unnested_length, total_length)?; - // Create the take indices array for other columns - let take_indices = create_take_indices(unnested_length, total_length); unnested_temp_arrays .into_iter() .zip(list_unnest_specs.iter()) @@ -715,7 +718,12 @@ fn list_unnest_at_level( // Dimension of arrays in batch is untouched, but the values are repeated // as the side effect of unnesting - let ret = repeat_arrs_from_indices(batch, &take_indices, &repeat_mask)?; + let ret = if repeat_mask.iter().any(|&repeat| repeat) { + repeat_arrs(batch, unnested_length, total_length, &repeat_mask)? + } else { + // Every column is a placeholder for an unnest result. No column consumes take indices. + batch.to_vec() + }; Ok(Some(ret)) } @@ -1247,11 +1255,11 @@ fn create_take_indices( PrimitiveArray::::from(indices) } -/// Create a batch of arrays based on an input `batch` and a `indices` array. -/// The `indices` array is used by the take kernel to repeat values in the arrays +/// Create a batch of arrays based on an input `batch` and per-row repetition lengths. +/// Take indices are created only for columns that need the take kernel to repeat values /// that are marked with `true` in the `repeat_mask`. Arrays marked with `false` -/// in the `repeat_mask` will be replaced with arrays filled with nulls of the -/// appropriate length. +/// are placeholders: they are replaced by their unnested versions before the +/// output batch is built. Keep those arrays instead of allocating unused null arrays. /// /// For example if we have the following batch: /// @@ -1267,11 +1275,10 @@ fn create_take_indices( /// c1: 1, null, 2, 3, 4, null, 5, 6 /// ``` /// -/// And the `indices` array contains the indices that are used by `take` kernel to -/// repeat the values in `c2`: +/// Each value in `c2` is repeated by its row's length: /// /// ```ignore -/// 0, 1, 2, 2, 2, 3, 4, 4 +/// 1, 1, 3, 1, 2 /// ``` /// /// so that the final batch will look like: @@ -1281,36 +1288,142 @@ fn create_take_indices( /// c2: 'a', 'b', 'c', 'c', 'c', null, 'd', 'd' /// ``` /// -/// The `repeat_mask` determines whether an array's values are repeated or replaced with nulls. -/// For example, if the `repeat_mask` is: -/// -/// ```ignore -/// [true, false] -/// ``` -/// -/// The final batch will look like: -/// -/// ```ignore -/// c1: 1, null, 2, 3, 4, null, 5, 6 // Repeated using `indices` -/// c2: null, null, null, null, null, null, null, null // Replaced with nulls -fn repeat_arrs_from_indices( +fn repeat_arrs( batch: &[ArrayRef], - indices: &PrimitiveArray, + lengths: &PrimitiveArray, + total_length: usize, repeat_mask: &[bool], ) -> Result>> { + let mut indices = None; batch .iter() .zip(repeat_mask.iter()) .map(|(arr, &repeat)| { if repeat { + if let Some(strings) = arr.as_string_opt::() { + if let Some(strings) = repeat_short_strings(strings, lengths, total_length) { + return Ok(Arc::new(strings) as ArrayRef); + } + } + let indices = + indices.get_or_insert_with(|| create_take_indices(lengths, total_length)); Ok(kernels::take::take(arr, indices, None)?) } else { - Ok(new_null_array(arr.data_type(), arr.len())) + Ok(Arc::clone(arr)) } }) .collect() } +/// Repeat short Utf8 values using the input runs rather than gathering every output value. +/// Batches dominated by long strings or with less than twofold expansion use Arrow's take instead. +fn repeat_short_strings( + array: &StringArray, + lengths: &PrimitiveArray, + total_length: usize, +) -> Option { + const SHORT_STRING_BYTES: usize = 16; + + debug_assert_eq!(array.len(), lengths.len()); + debug_assert_eq!(lengths.null_count(), 0); + debug_assert!(lengths.values().iter().all(|&n| n >= 0)); + if total_length < array.len().saturating_mul(2) { + return None; + } + if array.null_count() == array.len() { + return Some(StringArray::new_null(total_length)); + } + + // Reject wide input batches in O(1), before scanning their repetition lengths. This + // also avoids paying a guard pass followed by take for mixed short/long batches. + let input_offsets = array.value_offsets(); + let input_bytes = (input_offsets[array.len()] - input_offsets[0]) as usize; + if input_bytes > array.len().checked_mul(SHORT_STRING_BYTES)? { + return None; + } + let mut capacity = 0usize; + let mut valid_rows = 0usize; + for (i, (range, &repeat)) in input_offsets.windows(2).zip(lengths.values()).enumerate() { + if repeat == 0 || array.is_null(i) { + continue; + } + let len = (range[1] - range[0]) as usize; + capacity = capacity.checked_add(len.checked_mul(repeat as usize)?)?; + valid_rows = valid_rows.checked_add(repeat as usize)?; + // Stop early on long-string batches. Occasional longer values among short strings + // can use exact-size copies below, avoiding a full guard pass followed by take. + if capacity > valid_rows.checked_mul(SHORT_STRING_BYTES)? { + return None; + } + } + // Leave offset overflow to the existing take kernel, including its error reporting. + i32::try_from(capacity).ok()?; + let mut values = vec![0u8; capacity.checked_add(SHORT_STRING_BYTES)?]; + let mut offsets = Vec::with_capacity(total_length + 1); + offsets.push(0i32); + let mut nulls = (array.null_count() > 0).then(|| NullBufferBuilder::new(total_length)); + let data = array.value_data(); + let mut position = 0usize; + for (i, (range, &repeat)) in input_offsets.windows(2).zip(lengths.values()).enumerate() { + let repeat = repeat as usize; + if repeat == 0 { + continue; + } + if array.is_null(i) { + offsets.resize(offsets.len() + repeat, position as i32); + nulls.as_mut().unwrap().append_n_nulls(repeat); + continue; + } + let start = range[0] as usize; + let len = (range[1] - range[0]) as usize; + // Fixed-size copies avoid a tiny variable-size memcpy per output value. The source + // must contain the whole copy, even for a sliced array's final value. Extra bytes + // overwrite only future values or the output padding; later copies replace + // them before we truncate to the logical byte length. + if len == 0 { + offsets.resize(offsets.len() + repeat, position as i32); + } else if len <= 8 && data.len() - start >= 8 { + let value = &data[start..start + 8]; + for _ in 0..repeat { + values[position..position + 8].copy_from_slice(value); + position += len; + offsets.push(position as i32); + } + } else if len <= SHORT_STRING_BYTES && data.len() - start >= SHORT_STRING_BYTES { + let value = &data[start..start + SHORT_STRING_BYTES]; + for _ in 0..repeat { + values[position..position + SHORT_STRING_BYTES].copy_from_slice(value); + position += len; + offsets.push(position as i32); + } + } else { + let value = &data[start..start + len]; + for _ in 0..repeat { + values[position..position + len].copy_from_slice(value); + position += len; + offsets.push(position as i32); + } + } + if let Some(nulls) = &mut nulls { + nulls.append_n_non_nulls(repeat); + } + } + debug_assert_eq!(position, capacity); + debug_assert_eq!(offsets.len(), total_length + 1); + values.truncate(capacity); + let nulls = nulls + .and_then(|mut nulls| nulls.finish()) + .filter(|nulls| nulls.null_count() > 0); + // SAFETY: offsets are monotonic, fit i32, and end at the checked byte capacity. Each + // logical range copies a complete input Utf8 value, relying on the input array's Utf8 + // invariant; null ranges are empty. Safe slices bound every copy, values is fully + // initialized, and there is one validity entry per output value, by construction. + // Avoid validating the same input Utf8 again for every repetition. + Some(unsafe { + StringArray::new_unchecked(OffsetBuffer::new(offsets.into()), values.into(), nulls) + }) +} + #[cfg(test)] mod tests { use super::*; @@ -1325,6 +1438,208 @@ mod tests { use std::collections::VecDeque; use std::sync::atomic::{AtomicBool, Ordering as AtomicOrdering}; + #[test] + fn short_string_repetition_matches_take_buffers() { + let words = [ + "", + "a", + "abc", + "four", + "eight888", + "nine99999", + "sixteen_12345678", + "中文", + "😀", + "z", + ]; + for null_stride in [0usize, 2, 7, 1] { + let source = StringArray::from_iter((0..60).map(|i| { + (null_stride == 0 || i % null_stride != 0).then_some(words[i % words.len()]) + })); + for start in [0usize, 3, 59] { + let source = source.slice(start, 60 - start); + for fanout in [1i64, 5, 100] { + let lengths = Int64Array::from_iter_values((0..source.len()).map(|i| { + if i % 4 == 0 { + 0 + } else { + fanout + } + })); + let total = lengths.values().iter().sum::() as usize; + if total > source.len() { + assert_string_repeat_buffers(&source, &lengths); + } + } + } + } + // The last value has no following bytes to support a wider source copy. + for word in words { + assert_string_repeat_buffers( + &StringArray::from(vec![word]), + &Int64Array::from(vec![5]), + ); + } + // A null can own nonempty bytes; take must discard them. The middle slice also + // starts inside both the offset buffer and the validity bitmap. + let source = StringArray::new( + OffsetBuffer::new(vec![0, 3, 8, 9, 12].into()), + arrow::buffer::Buffer::from_slice_ref(b"abcHIDDENxyz"), + Some(NullBuffer::from(vec![true, false, false, true])), + ); + for lengths in [vec![5, 5, 5, 5], vec![5, 0, 0, 5], vec![0, 5, 5, 0]] { + assert_string_repeat_buffers(&source, &Int64Array::from(lengths)); + } + assert_string_repeat_buffers(&source.slice(1, 3), &Int64Array::from(vec![1, 0, 5])); + let valid = StringArray::new( + source.offsets().clone(), + source.values().clone(), + Some(NullBuffer::new_valid(source.len())), + ); + assert_string_repeat_buffers(&valid, &Int64Array::from(vec![5, 5, 5, 5])); + } + + fn assert_string_repeat_buffers(source: &StringArray, lengths: &Int64Array) { + let total = lengths.values().iter().sum::() as usize; + let actual = repeat_short_strings(source, lengths, total).expect("short string fast path"); + let indices = create_take_indices(lengths, total); + let expected = kernels::take::take(source, &indices, None).unwrap(); + let expected = expected.as_string::(); + assert_eq!(actual.value_offsets(), expected.value_offsets()); + assert_eq!(actual.value_data(), expected.value_data()); + assert_eq!(actual.nulls(), expected.nulls()); + actual.to_data().validate_full().unwrap(); + } + + #[test] + fn short_string_repetition_defers_long_unexpanded_and_overflowing_values() { + let lengths = Int64Array::from(vec![5]); + let long = StringArray::from(vec!["this string is longer than sixteen bytes"]); + assert!(repeat_short_strings(&long, &lengths, 5).is_none()); + let short = StringArray::from(vec!["abc"]); + assert!(repeat_short_strings(&short, &Int64Array::from(vec![1]), 1).is_none()); + assert!(repeat_short_strings(&short, &Int64Array::from(vec![0]), 0).is_none()); + let overflow = i32::MAX as i64; + assert!( + repeat_short_strings(&short, &Int64Array::from(vec![overflow]), overflow as usize) + .is_none() + ); + // Long values that are not selected do not disqualify the short values. + let mixed = StringArray::from(vec!["a very long unselected value", "abc"]); + assert_string_repeat_buffers(&mixed, &Int64Array::from(vec![0, 5])); + // A late long value among short strings must not trigger a complete guard pass + // followed by take. It uses an exact-size copy while the short values use wider ones. + let mixed = StringArray::from_iter_values((0..32).map(|i| { + if i == 31 { + "x".repeat(64) + } else { + "abc".to_owned() + } + })); + assert_string_repeat_buffers(&mixed, &Int64Array::from(vec![5; 32])); + let almost_flat = StringArray::from(vec!["a", "b", "c", "d"]); + assert!( + repeat_short_strings(&almost_flat, &Int64Array::from(vec![2, 1, 1, 1]), 5).is_none() + ); + } + + #[test] + fn unnest_placeholders_share_the_input() { + let list = Arc::clone(list_batch(&[Some(2), Some(1)]).column(0)); + let batch: Vec = vec![ + Arc::clone(&list), + Arc::new(Int32Array::from(vec![10, 20])), + Arc::new(StringArray::from(vec!["abc", "中文"])), + Arc::new(StringArray::from(vec![ + "a long value that uses Arrow take", + "another value longer than sixteen bytes", + ])), + Arc::new(Int64Array::from(vec![Some(1), None])), + ]; + let lengths = Int64Array::from(vec![2, 2]); + let result = repeat_arrs(&batch, &lengths, 4, &[false, true, true, true, true]).unwrap(); + assert!(Arc::ptr_eq(&result[0], &list)); + let indices = create_take_indices(&lengths, 4); + for (input, output) in batch.iter().zip(&result).skip(1) { + let expected = kernels::take::take(input, &indices, None).unwrap(); + assert_eq!(output.to_data(), expected.to_data()); + } + } + + #[test] + fn recursive_unnest_replaces_placeholders_with_and_without_carried_columns() { + let inner = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(10), None]), + None, + Some(vec![]), + Some(vec![Some(20)]), + Some(vec![Some(30), Some(40)]), + ]); + let outer: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new("item", inner.data_type().clone(), true)), + OffsetBuffer::new(ScalarBuffer::from(vec![0, 3, 5, 5, 5])), + Arc::new(inner), + Some(NullBuffer::from(vec![true, true, false, true])), + )); + for carried in [false, true] { + for handling in [NullHandling::Drop, NullHandling::PreserveAndExpandEmpty] { + let mut columns = vec![Arc::clone(&outer)]; + let mut fields = vec![Field::new("lists", outer.data_type().clone(), true)]; + let mut output_fields = vec![Field::new("value", DataType::Int32, true)]; + if carried { + columns.push(Arc::new(Int32Array::from(vec![1, 2, 3, 4]))); + fields.push(Field::new("id", DataType::Int32, false)); + output_fields.push(Field::new("id", DataType::Int32, false)); + } + let input = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap(); + let output = build_batch( + &input, + &Arc::new(Schema::new(output_fields)), + &[ListUnnest { + index_in_input_schema: 0, + depth: 2, + }], + &HashSet::new(), + &UnnestOptions::new().with_null_handling(handling), + None, + ) + .unwrap() + .unwrap(); + let (expected, ids) = if handling == NullHandling::Drop { + ( + vec![Some(10), None, Some(20), Some(30), Some(40)], + vec![1, 1, 2, 2, 2], + ) + } else { + ( + vec![ + Some(10), + None, + None, + None, + Some(20), + Some(30), + Some(40), + None, + None, + ], + vec![1, 1, 1, 1, 2, 2, 2, 3, 4], + ) + }; + assert_eq!( + output.column(0).as_primitive::(), + &Int32Array::from(expected) + ); + if carried { + assert_eq!( + output.column(1).as_primitive::(), + &Int32Array::from(ids) + ); + } + } + } + } + /// Build a single-column `List` batch where row `i` holds `lens[i]` elements, /// numbered consecutively across the whole batch. `None` is a NULL list. fn list_batch(lens: &[Option]) -> RecordBatch { diff --git a/native/proto/src/proto/expr.proto b/native/proto/src/proto/expr.proto index 096773af7e5..b85f08f6531 100644 --- a/native/proto/src/proto/expr.proto +++ b/native/proto/src/proto/expr.proto @@ -434,7 +434,8 @@ message Cast { DataType datatype = 2; string timezone = 3; EvalMode eval_mode = 4; - bool allow_incompat = 5; + reserved 5; + reserved "allow_incompat"; // True when running against Spark 4.0+. Controls version-specific cast behaviour // such as the handling of leading whitespace before T-prefixed time-only strings. bool is_spark4_plus = 6; @@ -597,6 +598,11 @@ message TruncTimestamp { Expr format = 1; Expr child = 2; string timezone = 3; + // Whether SECOND and MILLISECOND truncation wraps when the result would fall below the + // smallest timestamp. Spark before 4.2.0 wraps the Long subtraction; 4.2.0 and later + // (SPARK-56663) raise `long overflow` like every other unit. This tracks the Spark version Comet is + // running against. + bool wrap_second_millisecond_overflow = 4; } message Subquery { diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 527461f47a5..8225bad35c0 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -785,7 +785,7 @@ message IcebergWriteCommon { } // Single Iceberg write operator. Per-task fields are populated by the JVM exec wrapper at task -// launch (matching how `ParquetWriter.task_attempt_id` is filled in `CometNativeWriteExec`). +// launch (matching how `ParquetWriter.output_path` is filled in by the Parquet write execs). message IcebergWrite { IcebergWriteCommon common = 1; @@ -920,27 +920,19 @@ message ShuffleWriter { } message ParquetWriter { - // Where this task writes. Two shapes, selected by whether `work_dir` is set: - // - // - Spark 4.0+ (CometWriteFilesExec): `work_dir` is unset and `output_path` is the - // fully-qualified path of the Parquet file to write, set per task from - // FileCommitProtocol.newTaskTempFile. Naming and staging are owned by Spark's commit - // protocol so that task-attempt isolation, speculative execution, and committers that track - // individual files (S3A magic, streaming manifest) all behave as they do for Spark's own - // writer. The native writer uses this path verbatim. - // - Spark 3.x (CometNativeWriteExec): `work_dir` is set and the native writer derives the file - // name from it, the partition id and the task attempt id. `output_path` is the write's - // output directory and is unused natively. Goes away with Spark 3.x support. + // The fully-qualified path of the Parquet file this task writes, set per task from + // FileCommitProtocol.newTaskTempFile by CometWriteFilesExec (Spark 4.0+) or + // CometNativeWriteExec (Spark 3.x). Naming and staging are owned by Spark's commit protocol so + // that task-attempt isolation, speculative execution, and committers that track individual + // files (S3A magic, streaming manifest) all behave as they do for Spark's own writer. The native + // writer uses this path verbatim. string output_path = 1; CompressionCodec compression = 2; repeated string column_names = 4; - // Working directory for temporary files (used by FileCommitProtocol). Spark 3.x only; see - // output_path above. - optional string work_dir = 5; - // Job ID for tracking this write operation - optional string job_id = 6; - // Task attempt ID for this specific task - optional int32 task_attempt_id = 7; + // Formerly work_dir, job_id and task_attempt_id, which let the native writer name its own file + // under a working directory. + reserved 5, 6, 7; + reserved "work_dir", "job_id", "task_attempt_id"; // Options for configuring object stores such as AWS S3, GCS, etc. The key-value pairs are taken // from Hadoop configuration for compatibility with Hadoop FileSystem implementations of object // stores. @@ -1012,9 +1004,9 @@ message MergeOutputRow { } message Explode { - // The array expression to explode into multiple rows + // The array or map expression to explode into multiple rows spark.spark_expression.Expr child = 1; - // Whether this is explode_outer (produces null row for empty/null arrays) + // Whether this is explode_outer (produces null row for empty/null collections) bool outer = 2; // Expressions for other columns to project alongside the exploded values repeated spark.spark_expression.Expr project_list = 3; diff --git a/native/shuffle/src/spark_unsafe/list.rs b/native/shuffle/src/spark_unsafe/list.rs index 3d8e89ebc56..f74fcf95baa 100644 --- a/native/shuffle/src/spark_unsafe/list.rs +++ b/native/shuffle/src/spark_unsafe/list.rs @@ -27,10 +27,15 @@ use arrow::array::{ ListBuilder, NullBuilder, StringBuilder, StructBuilder, Time64NanosecondBuilder, TimestampMicrosecondBuilder, }, - MapBuilder, + MapBuilder, PrimitiveArray, +}; +use arrow::buffer::{BooleanBuffer, Buffer, NullBuffer, ScalarBuffer}; +use arrow::datatypes::{ + DataType, Date32Type, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type, + Time64NanosecondType, TimeUnit, TimestampMicrosecondType, }; -use arrow::datatypes::{DataType, TimeUnit}; use datafusion_comet_jni_bridge::errors::CometError; +use std::sync::Arc; /// Element count from which a nullable array with no null elements is appended with one copy /// instead of element by element. Below it the loop is as fast: checking the null bitset and @@ -39,55 +44,93 @@ use datafusion_comet_jni_bridge::errors::CometError; /// time with the copy, and arrays of 50 about a fifth. const MIN_BULK_APPEND_ELEMENTS: usize = 8; +/// Element count from which an aligned nullable array that holds a null is appended with one copy +/// of its values and a validity buffer built from its null bitset, instead of element by element. +/// Building the two buffers takes four allocations, which cost as much as appending about 50 +/// elements one by one, so shorter arrays stay on the loop. The `list_with_one_null` group of the +/// `array_element_append` benchmark times lengths on both sides of the cutoff, and its doc comment +/// shows how to compare the two paths. On an M3 Max with the system allocator, arrays of 32 +/// elements took 1.4 times as long with the copy, and arrays of 64 took 69% of the time for 4-byte +/// elements and 86% for 8-byte ones. +const MIN_BULK_NULLABLE_APPEND_ELEMENTS: usize = 64; + /// Generates bulk append methods for primitive types in SparkUnsafeArray. /// /// # Safety invariants for all generated methods: /// - `element_offset` points to contiguous element data of length `num_elements` /// - `null_bitset_ptr()` returns a pointer to `ceil(num_elements/64)` i64 words /// - These invariants are guaranteed by the SparkUnsafeArray layout from the JVM +/// +/// A timestamp appender also takes `timezone`, which must be the builder's timezone, because +/// `append_array` requires the array's data type to match the builder's. A mismatch panics in +/// `append_array` for an aligned array of `MIN_BULK_NULLABLE_APPEND_ELEMENTS` or more elements that +/// holds a null, and goes unnoticed for every other array. macro_rules! impl_append_to_builder { - ($method_name:ident, $builder_type:ty, $element_type:ty) => { - pub(crate) fn $method_name(&self, builder: &mut $builder_type) { + ($method_name:ident, $builder_type:ty, $element_type:ty, $arrow_type:ty + $(, $timezone:ident)?) => { + pub(crate) fn $method_name( + &self, + builder: &mut $builder_type, + $($timezone: &Option>,)? + ) { let num_elements = self.num_elements; if num_elements == 0 { return; } + debug_assert!(self.element_offset != 0, "element_offset is null"); + let ptr = self.element_offset as *const $element_type; + // An array can start at an address that is not aligned for its elements (see + // `SparkUnsafeObject`), and only aligned elements can be read as a slice. + let aligned = (ptr as usize).is_multiple_of(std::mem::align_of::<$element_type>()); // A nullable array without nulls takes the copy below too, once it is long enough. if NULLABLE && (num_elements < MIN_BULK_APPEND_ELEMENTS || self.has_null()) { - let mut ptr = self.element_offset as *const $element_type; let null_words = self.null_bitset_ptr(); - debug_assert!(!null_words.is_null(), "null_bitset_ptr is null"); - debug_assert!(!ptr.is_null(), "element_offset pointer is null"); - for idx in 0..num_elements { - // SAFETY: null_words has ceil(num_elements/64) words, idx < num_elements - let is_null = unsafe { Self::is_null_in_bitset(null_words, idx) }; - - if is_null { - builder.append_null(); - } else { - // SAFETY: ptr is within element data bounds - builder.append_value(unsafe { ptr.read_unaligned() }); - } - // SAFETY: ptr stays within bounds, iterating num_elements times - ptr = unsafe { ptr.add(1) }; - } - } else { - // SAFETY: element_offset points to contiguous data of length num_elements - debug_assert!(self.element_offset != 0, "element_offset is null"); - let ptr = self.element_offset as *const $element_type; - // Use bulk copy when data is properly aligned, fall back to - // per-element unaligned reads otherwise - if (ptr as usize).is_multiple_of(std::mem::align_of::<$element_type>()) { - let slice = unsafe { std::slice::from_raw_parts(ptr, num_elements) }; - builder.append_slice(slice); + if aligned && num_elements >= MIN_BULK_NULLABLE_APPEND_ELEMENTS { + // SAFETY: element_offset points to num_elements aligned elements. + let values = unsafe { std::slice::from_raw_parts(ptr, num_elements) }; + let null_mask_len = num_elements.div_ceil(8); + // SAFETY: the ceil(num_elements/64) words of the null bitset cover these + // bytes. The words are little-endian, so byte i holds the bits of elements + // 8i to 8i+7, least significant first, as in an Arrow bitmap. + let null_mask = unsafe { + std::slice::from_raw_parts(null_words as *const u8, null_mask_len) + }; + // Spark sets the bit of a null element and Arrow the bit of a valid one. The + // padding bits past the last element turn to ones, outside the buffer's length. + let flipped: Vec = null_mask.iter().map(|byte| !byte).collect(); + let validity = + NullBuffer::new(BooleanBuffer::new(Buffer::from(flipped), 0, num_elements)); + let arr = PrimitiveArray::<$arrow_type>::new( + ScalarBuffer::from(Buffer::from_slice_ref(values)), + Some(validity), + ); + $(let arr = arr.with_timezone_opt($timezone.clone());)? + builder.append_array(&arr); } else { let mut ptr = ptr; - for _ in 0..num_elements { - builder.append_value(unsafe { ptr.read_unaligned() }); + for idx in 0..num_elements { + // SAFETY: null_words has ceil(num_elements/64) words, idx < num_elements + if unsafe { Self::is_null_in_bitset(null_words, idx) } { + builder.append_null(); + } else { + // SAFETY: ptr is within element data bounds + builder.append_value(unsafe { ptr.read_unaligned() }); + } + // SAFETY: ptr stays within bounds, iterating num_elements times ptr = unsafe { ptr.add(1) }; } } + } else if aligned { + // SAFETY: element_offset points to num_elements aligned elements. + let slice = unsafe { std::slice::from_raw_parts(ptr, num_elements) }; + builder.append_slice(slice); + } else { + let mut ptr = ptr; + for _ in 0..num_elements { + builder.append_value(unsafe { ptr.read_unaligned() }); + ptr = unsafe { ptr.add(1) }; + } } } }; @@ -180,7 +223,9 @@ impl SparkUnsafeArray { fn has_null(&self) -> bool { let null_words = self.null_bitset_ptr(); // SAFETY: the null bitset holds ceil(num_elements/64) words. Spark zeroes the bits past - // the last element, and a stray one would only send the array down the per-element path. + // the last element. A stray one would send the array to the per-element loop, or to the + // validity buffer if it is aligned and long enough, and both ignore the bits past the + // last element. (0..self.num_elements.div_ceil(64)) .any(|word| unsafe { null_words.add(word).read_unaligned() } != 0) } @@ -196,19 +241,26 @@ impl SparkUnsafeArray { (null_words.add(word_idx).read_unaligned() & (1i64 << bit_idx)) != 0 } - impl_append_to_builder!(append_ints_to_builder, Int32Builder, i32); - impl_append_to_builder!(append_longs_to_builder, Int64Builder, i64); - impl_append_to_builder!(append_time64s_to_builder, Time64NanosecondBuilder, i64); - impl_append_to_builder!(append_shorts_to_builder, Int16Builder, i16); - impl_append_to_builder!(append_bytes_to_builder, Int8Builder, i8); - impl_append_to_builder!(append_floats_to_builder, Float32Builder, f32); - impl_append_to_builder!(append_doubles_to_builder, Float64Builder, f64); + impl_append_to_builder!(append_ints_to_builder, Int32Builder, i32, Int32Type); + impl_append_to_builder!(append_longs_to_builder, Int64Builder, i64, Int64Type); + impl_append_to_builder!( + append_time64s_to_builder, + Time64NanosecondBuilder, + i64, + Time64NanosecondType + ); + impl_append_to_builder!(append_shorts_to_builder, Int16Builder, i16, Int16Type); + impl_append_to_builder!(append_bytes_to_builder, Int8Builder, i8, Int8Type); + impl_append_to_builder!(append_floats_to_builder, Float32Builder, f32, Float32Type); + impl_append_to_builder!(append_doubles_to_builder, Float64Builder, f64, Float64Type); impl_append_to_builder!( append_timestamps_to_builder, TimestampMicrosecondBuilder, - i64 + i64, + TimestampMicrosecondType, + timezone ); - impl_append_to_builder!(append_dates_to_builder, Date32Builder, i32); + impl_append_to_builder!(append_dates_to_builder, Date32Builder, i32, Date32Type); /// Bulk append boolean values to builder. /// Booleans are stored as 1 byte each in SparkUnsafeArray, requiring special handling. @@ -220,37 +272,28 @@ impl SparkUnsafeArray { if num_elements == 0 { return; } - - let mut ptr = self.element_offset as *const u8; debug_assert!( - !ptr.is_null(), + self.element_offset != 0, "append_booleans: element_offset pointer is null" ); + // SAFETY: element_offset points to num_elements one-byte booleans, which need no + // alignment. + let values = + unsafe { std::slice::from_raw_parts(self.element_offset as *const u8, num_elements) }; if NULLABLE { let null_words = self.null_bitset_ptr(); - debug_assert!( - !null_words.is_null(), - "append_booleans: null_bitset_ptr is null" - ); - for idx in 0..num_elements { + for (idx, &value) in values.iter().enumerate() { // SAFETY: null_words has ceil(num_elements/64) words, idx < num_elements - let is_null = unsafe { Self::is_null_in_bitset(null_words, idx) }; - - if is_null { + if unsafe { Self::is_null_in_bitset(null_words, idx) } { builder.append_null(); } else { - // SAFETY: ptr is within element data bounds - builder.append_value(unsafe { *ptr != 0 }); + builder.append_value(value != 0); } - // SAFETY: ptr stays within bounds, iterating num_elements times - ptr = unsafe { ptr.add(1) }; } } else { - for _ in 0..num_elements { - // SAFETY: ptr is within element data bounds - builder.append_value(unsafe { *ptr != 0 }); - ptr = unsafe { ptr.add(1) }; + for &value in values { + builder.append_value(value != 0); } } } @@ -307,9 +350,9 @@ pub fn append_to_builder( let builder = downcast_builder_ref!(Float64Builder, builder); array.append_doubles_to_builder::(builder); } - DataType::Timestamp(TimeUnit::Microsecond, _) => { + DataType::Timestamp(TimeUnit::Microsecond, tz) => { let builder = downcast_builder_ref!(TimestampMicrosecondBuilder, builder); - array.append_timestamps_to_builder::(builder); + array.append_timestamps_to_builder::(builder, tz); } DataType::Date32 => { let builder = downcast_builder_ref!(Date32Builder, builder); @@ -415,6 +458,7 @@ pub fn append_list_element( mod test { use super::*; use arrow::array::builder::PrimitiveBuilder; + use arrow::array::Array; use arrow::datatypes::{ ArrowPrimitiveType, Date32Type, Int16Type, Int32Type, Int64Type, Int8Type, TimestampMicrosecondType, @@ -455,7 +499,7 @@ mod test { let base = unsafe { (words.as_mut_ptr() as *mut u8).add(shift) }; unsafe { std::ptr::copy_nonoverlapping(bytes.as_ptr(), base, bytes.len()) }; let array = SparkUnsafeArray::new(base as i64); - let mut builder = PrimitiveBuilder::::new(); + let mut builder = PrimitiveBuilder::::new().with_data_type(data_type.clone()); append_to_builder::(&data_type, &mut builder, &array).unwrap(); builder.finish().iter().map(|v| v.map(Into::into)).collect() } @@ -465,6 +509,11 @@ mod test { let long: Vec> = (0..70).map(Some).collect(); let mut null_in_second_bitset_word = long.clone(); null_in_second_bitset_word[66] = None; + let nulls_at_both_ends = |len: i64| -> Vec> { + (0..len) + .map(|v| (v != 0 && v != len - 1).then_some(-v)) + .collect() + }; let cases = [ vec![], // Shorter than MIN_BULK_APPEND_ELEMENTS: element by element. @@ -472,8 +521,22 @@ mod test { // Holds a null: element by element. vec![Some(1), None, Some(3), Some(-4), Some(5)], // No nulls and long enough: one copy. - vec![Some(-1), Some(2), Some(-3), Some(4)], + vec![ + Some(-1), + Some(2), + Some(-3), + Some(4), + Some(-5), + Some(6), + Some(-7), + Some(8), + ], long, + // Holds a null, one element short of MIN_BULK_NULLABLE_APPEND_ELEMENTS: element by + // element. + nulls_at_both_ends(63), + // Holds a null and long enough: one copy and a validity buffer. + nulls_at_both_ends(64), // The null is only in the second bitset word, which the copy check must read. null_in_second_bitset_word, ]; @@ -508,7 +571,116 @@ mod test { ), *values ); + // The builder carries the column's timezone, which `append_array` checks. + assert_eq!( + round_trip::( + DataType::Timestamp(TimeUnit::Microsecond, Some("America/Denver".into())), + values, + shift + ), + *values + ); } } } + + fn make_i32_array(num_elements: usize, null_indices: &[usize]) -> Vec { + let null_bitset_words = num_elements.div_ceil(64); + let header_size = 8 + null_bitset_words * 8; + let mut buffer = vec![0u8; header_size + num_elements * 4]; + + buffer[0..8].copy_from_slice(&(num_elements as i64).to_le_bytes()); + + for &i in null_indices { + let word_offset = 8 + (i / 64) * 8; + let word = i64::from_le_bytes(buffer[word_offset..word_offset + 8].try_into().unwrap()); + buffer[word_offset..word_offset + 8] + .copy_from_slice(&(word | (1i64 << (i % 64))).to_le_bytes()); + } + + for i in 0..num_elements { + let offset = header_size + i * 4; + buffer[offset..offset + 4].copy_from_slice(&(i as i32).to_le_bytes()); + } + buffer + } + + /// Copies `bytes` to an 8-byte boundary, as inside an UnsafeRow. + fn aligned_copy(bytes: &[u8]) -> Vec { + let mut words = vec![0u64; bytes.len().div_ceil(8)]; + // SAFETY: `words` covers `bytes.len()` bytes. + unsafe { + std::ptr::copy_nonoverlapping( + bytes.as_ptr(), + words.as_mut_ptr() as *mut u8, + bytes.len(), + ) + }; + words + } + + fn append_i32(buffer: &[u8]) -> arrow::array::Int32Array { + let words = aligned_copy(buffer); + let array = SparkUnsafeArray::new(words.as_ptr() as i64); + let mut builder = Int32Builder::new(); + append_to_builder::(&DataType::Int32, &mut builder, &array).unwrap(); + builder.finish() + } + + #[test] + fn test_nullable_null_in_last_byte() { + // None of the lengths is a multiple of 8, so the null sits in a partial last byte. The + // first three take the per-element loop, the rest the validity buffer. + for num_elements in [7, 9, 17, 65, 71, 73, 81] { + let last = num_elements - 1; + let arr = append_i32(&make_i32_array(num_elements, &[last])); + assert_eq!(arr.len(), num_elements); + assert_eq!(arr.null_count(), 1); + assert!(arr.is_null(last)); + for i in 0..last { + assert!(arr.is_valid(i)); + assert_eq!(arr.value(i), i as i32); + } + } + } + + #[test] + fn test_all_null() { + for num_elements in [9, 73] { + let indices: Vec = (0..num_elements).collect(); + let arr = append_i32(&make_i32_array(num_elements, &indices)); + assert_eq!(arr.len(), num_elements); + assert_eq!(arr.null_count(), num_elements); + } + } + + #[test] + fn test_all_valid() { + for num_elements in [9, 73] { + let arr = append_i32(&make_i32_array(num_elements, &[])); + assert_eq!(arr.len(), num_elements); + assert_eq!(arr.null_count(), 0); + } + } + + #[test] + fn nullable_arrays_append_into_one_builder() { + // The list path appends every row's array to one values builder, so the validity buffer + // of a long array is appended at an offset that is not a whole byte. + let arrays: [Vec>; 4] = [ + vec![Some(1), None, Some(3)], + (0..70).map(|v| (v % 9 != 4).then_some(v)).collect(), + vec![None; 5], + (0..65).map(|v| (v != 64).then_some(-v)).collect(), + ]; + let mut builder = Int32Builder::new(); + for values in &arrays { + let words = aligned_copy(&unsafe_array_bytes(values, 4)); + let array = SparkUnsafeArray::new(words.as_ptr() as i64); + append_to_builder::(&DataType::Int32, &mut builder, &array).unwrap(); + } + let appended: Vec> = + builder.finish().iter().map(|v| v.map(i64::from)).collect(); + assert_eq!(appended, arrays.concat()); + } } diff --git a/native/shuffle/src/spark_unsafe/row.rs b/native/shuffle/src/spark_unsafe/row.rs index 180673c37fb..ce18a398df5 100644 --- a/native/shuffle/src/spark_unsafe/row.rs +++ b/native/shuffle/src/spark_unsafe/row.rs @@ -622,7 +622,7 @@ fn append_list_column_batch( // Helper macro for primitive element types - gets builder fresh each iteration // to avoid borrow conflicts with list_builder.append() macro_rules! process_primitive_lists { - ($builder_type:ty, $append_fn:ident) => {{ + ($builder_type:ty, $append_fn:ident $(, $extra:expr)*) => {{ for i in row_start..row_end { read_row_at!(row, row_addresses_ptr, row_sizes_ptr, i); @@ -636,7 +636,7 @@ fn append_list_column_batch( .as_any_mut() .downcast_mut::<$builder_type>() .expect(stringify!($builder_type)); - array.$append_fn::(values_builder); + array.$append_fn::(values_builder $(, $extra)*); list_builder.append(true); } } @@ -668,8 +668,12 @@ fn append_list_column_batch( DataType::Date32 => { process_primitive_lists!(Date32Builder, append_dates_to_builder); } - DataType::Timestamp(TimeUnit::Microsecond, _) => { - process_primitive_lists!(TimestampMicrosecondBuilder, append_timestamps_to_builder); + DataType::Timestamp(TimeUnit::Microsecond, tz) => { + process_primitive_lists!( + TimestampMicrosecondBuilder, + append_timestamps_to_builder, + tz + ); } DataType::Time64(TimeUnit::Nanosecond) => { process_primitive_lists!(Time64NanosecondBuilder, append_time64s_to_builder); @@ -1504,7 +1508,8 @@ fn make_batch(arrays: Vec, row_count: usize) -> Result`. assert_eq!(&*row.get_string(0), "\u{FFFD}\u{FFFD}A"); } + + /// Lays out an 8-byte-aligned `UnsafeRow` with one array field of 8-byte elements: the null + /// bitset, the field's offset and size, then the `UnsafeArrayData`. `None` is a null field. + fn unsafe_row_with_array(values: Option<&[Option]>) -> Vec { + let Some(values) = values else { + return vec![1, 0]; + }; + let mut null_bitset = vec![0u64; values.len().div_ceil(64)]; + let mut elements = Vec::with_capacity(values.len()); + for (i, value) in values.iter().enumerate() { + if value.is_none() { + null_bitset[i / 64] |= 1 << (i % 64); + } + elements.push(value.unwrap_or(0) as u64); + } + let array_size = 8 * (1 + null_bitset.len() + elements.len()) as u64; + let mut row = vec![0, (16 << 32) | array_size, values.len() as u64]; + row.extend(null_bitset); + row.extend(elements); + row + } + + #[test] + fn timestamp_list_column_keeps_its_timezone() { + // An array of at least MIN_BULK_NULLABLE_APPEND_ELEMENTS elements that holds a null is + // appended with `append_array`, which panics unless the array has the builder's timezone. + let data_type = DataType::List(Arc::new(Field::new( + "item", + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + true, + ))); + let with_null: Vec> = (0..70).map(|v| (v != 3).then_some(v << 20)).collect(); + let without_null: Vec> = (0..70).map(|v| Some(-v)).collect(); + let short = vec![Some(1), None]; + let lists = [ + Some(with_null.as_slice()), + None, + Some(short.as_slice()), + Some(without_null.as_slice()), + ]; + let rows: Vec> = lists.iter().map(|l| unsafe_row_with_array(*l)).collect(); + let mut row_addresses: Vec = rows.iter().map(|r| r.as_ptr() as jlong).collect(); + let mut row_sizes: Vec = rows.iter().map(|r| (r.len() * 8) as jint).collect(); + + let schema = [data_type.clone()]; + let mut builder = make_builders(&data_type, rows.len(), 1.0).unwrap(); + append_columns( + row_addresses.as_mut_ptr(), + row_sizes.as_mut_ptr(), + 0, + rows.len(), + &schema, + 0, + &mut builder, + 1.0, + ) + .unwrap(); + let array = builder.finish(); + + assert_eq!(array.data_type(), &data_type); + let appended: Vec>>> = array + .as_list::() + .iter() + .map(|values| { + values.map(|values| { + values + .as_primitive::() + .iter() + .collect() + }) + }) + .collect(); + let expected: Vec>>> = + lists.iter().map(|l| l.map(<[_]>::to_vec)).collect(); + assert_eq!(appended, expected); + } } diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index d177a42fb8c..3011f1c5c6f 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -30,6 +30,7 @@ edition = { workspace = true } arrow = { workspace = true } chrono = { workspace = true } datafusion = { workspace = true } +datafusion-functions = { workspace = true } datafusion-spark = { workspace = true } num = { workspace = true } regex = { workspace = true } @@ -82,6 +83,10 @@ harness = false name = "cast_from_string" harness = false +[[bench]] +name = "cast_nested" +harness = false + [[bench]] name = "cast_string_to_date" harness = false @@ -182,6 +187,10 @@ harness = false name = "bloom_filter_agg" harness = false +[[bench]] +name = "bloom_filter_might_contain" +harness = false + [[bench]] name = "collect" harness = false @@ -454,6 +463,14 @@ harness = false name = "aggregate_stats" harness = false +[[bench]] +name = "hll_agg" +harness = false + +[[bench]] +name = "hll_scalar" +harness = false + [[bench]] name = "aggregate_misc" harness = false @@ -474,6 +491,10 @@ harness = false name = "nested_comparison" harness = false +[[bench]] +name = "float_comparison" +harness = false + [[bench]] name = "float_arrays" harness = false diff --git a/native/spark-expr/benches/bloom_filter_might_contain.rs b/native/spark-expr/benches/bloom_filter_might_contain.rs new file mode 100644 index 00000000000..c0e926a40ff --- /dev/null +++ b/native/spark-expr/benches/bloom_filter_might_contain.rs @@ -0,0 +1,144 @@ +// 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. + +//! Benchmark for the per-row `might_contain` probe (`BloomFilterMightContain`). The filter it +//! probes against is built up front by running the public `BloomFilterAgg` over a column of +//! longs, so the bench exercises only the probe path (`might_contain_longs`) and not filter +//! construction. + +use arrow::array::{Array, ArrayRef, BinaryArray, Int64Array, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema}; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::common::config::ConfigOptions; +use datafusion::common::ScalarValue; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; +use datafusion::logical_expr::{AggregateUDF, ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; +use datafusion::physical_expr::aggregate::AggregateExprBuilder; +use datafusion::physical_expr::expressions::{lit, Column}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; +use datafusion::physical_plan::{collect, ExecutionPlan}; +use datafusion::prelude::SessionContext; +use datafusion_comet_spark_expr::{ + BloomFilterAgg, BloomFilterMightContain, SparkBloomFilterVersion, +}; +use std::hint::black_box; +use std::sync::Arc; +use tokio::runtime::Runtime; + +const DISTINCT_ITEMS: usize = 8192; +const NUM_BITS: i32 = (DISTINCT_ITEMS as i32) * 16; + +/// Build a serialized Spark bloom filter containing `0..DISTINCT_ITEMS` by running the public +/// `BloomFilterAgg` once. Returns the filter's Spark-serialized bytes. +fn build_filter_bytes(rt: &Runtime) -> Vec { + let values: ArrayRef = Arc::new(Int64Array::from( + (0..DISTINCT_ITEMS as i64).collect::>(), + )); + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![values]).unwrap(); + + let scan: Arc = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], Arc::clone(&schema), None).unwrap(), + ))); + let udf = AggregateUDF::new_from_impl(BloomFilterAgg::new( + lit(ScalarValue::Int64(Some(DISTINCT_ITEMS as i64))), + lit(ScalarValue::Int64(Some(NUM_BITS as i64))), + DataType::Binary, + SparkBloomFilterVersion::V1, + )); + let aggr_expr = AggregateExprBuilder::new( + Arc::new(udf), + vec![Arc::new(Column::new("v", 0)) as Arc], + ) + .schema(Arc::clone(&schema)) + .alias("bloom_filter") + .build() + .unwrap(); + + let aggregate: Arc = Arc::new( + AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new(vec![], vec![], vec![vec![]], false), + vec![aggr_expr.into()], + vec![None], + scan, + schema, + ) + .unwrap(), + ); + + let ctx = SessionContext::new(); + let batches = rt.block_on(collect(aggregate, ctx.task_ctx())).unwrap(); + let col = batches[0].column(0); + let bytes = col.as_any().downcast_ref::().unwrap(); + bytes.value(0).to_vec() +} + +/// Build an Int64 probe column: the first half are hits (`0..rows/2`, present in the filter), +/// the second half are misses (large values not inserted). Every `null_every`-th row is null +/// (`null_every == 0` means no nulls). +fn create_probe_array(rows: usize, null_every: usize) -> ArrayRef { + let arr: Int64Array = (0..rows) + .map(|i| { + if null_every != 0 && i % null_every == 0 { + None + } else if i < rows / 2 { + Some(i as i64) + } else { + Some(i as i64 + 1_000_000_000) + } + }) + .collect(); + Arc::new(arr) +} + +fn run(udf: &BloomFilterMightContain, values: &ArrayRef, rows: usize) { + let args = vec![ColumnarValue::Array(Arc::clone(values))]; + black_box( + udf.invoke_with_args(ScalarFunctionArgs { + args, + arg_fields: vec![], + number_rows: rows, + return_field: Arc::new(Field::new("result", DataType::Boolean, true)), + config_options: Arc::new(ConfigOptions::default()), + }) + .unwrap(), + ); +} + +fn criterion_benchmark(c: &mut Criterion) { + let rows = 8192; + let rt = Runtime::new().unwrap(); + let filter_bytes = build_filter_bytes(&rt); + let udf = + BloomFilterMightContain::try_new(lit(ScalarValue::Binary(Some(filter_bytes)))).unwrap(); + + let no_nulls = create_probe_array(rows, 0); + let with_nulls = create_probe_array(rows, 10); + + let mut group = c.benchmark_group("bloom_filter_might_contain"); + group.bench_function("probe no nulls", |b| b.iter(|| run(&udf, &no_nulls, rows))); + group.bench_function("probe with nulls", |b| { + b.iter(|| run(&udf, &with_nulls, rows)) + }); + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/benches/cast_binary_to_string.rs b/native/spark-expr/benches/cast_binary_to_string.rs index 8ae9b617615..7c3e5d3a9c5 100644 --- a/native/spark-expr/benches/cast_binary_to_string.rs +++ b/native/spark-expr/benches/cast_binary_to_string.rs @@ -46,7 +46,7 @@ fn build(size: usize, width: usize, null_every: usize) -> ArrayRef { } fn options(style: Option) -> SparkCastOptions { - let mut options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkCastOptions::new(EvalMode::Legacy, "UTC"); options.binary_output_style = style; options } diff --git a/native/spark-expr/benches/cast_decimal_to_boolean.rs b/native/spark-expr/benches/cast_decimal_to_boolean.rs index 0d1bfe255c3..2e5af4809e2 100644 --- a/native/spark-expr/benches/cast_decimal_to_boolean.rs +++ b/native/spark-expr/benches/cast_decimal_to_boolean.rs @@ -59,7 +59,7 @@ fn criterion_benchmark(c: &mut Criterion) { let cast_to_bool = Cast::new( expr, DataType::Boolean, - SparkCastOptions::new(EvalMode::Legacy, "UTC", false), + SparkCastOptions::new(EvalMode::Legacy, "UTC"), None, None, ); diff --git a/native/spark-expr/benches/cast_decimal_to_string.rs b/native/spark-expr/benches/cast_decimal_to_string.rs index d0629300a97..5fd76e44fd4 100644 --- a/native/spark-expr/benches/cast_decimal_to_string.rs +++ b/native/spark-expr/benches/cast_decimal_to_string.rs @@ -43,7 +43,7 @@ fn cast_to_utf8() -> Cast { Cast::new( Arc::new(Column::new("a", 0)), DataType::Utf8, - SparkCastOptions::new_without_timezone(EvalMode::Legacy, false), + SparkCastOptions::new_without_timezone(EvalMode::Legacy), None, None, ) diff --git a/native/spark-expr/benches/cast_float_to_decimal.rs b/native/spark-expr/benches/cast_float_to_decimal.rs index fab3c3d1122..1bb071092d3 100644 --- a/native/spark-expr/benches/cast_float_to_decimal.rs +++ b/native/spark-expr/benches/cast_float_to_decimal.rs @@ -52,7 +52,7 @@ fn cast(to: DataType, mode: EvalMode) -> Cast { Cast::new( Arc::new(Column::new("a", 0)), to, - SparkCastOptions::new_without_timezone(mode, false), + SparkCastOptions::new_without_timezone(mode), None, None, ) diff --git a/native/spark-expr/benches/cast_float_to_string.rs b/native/spark-expr/benches/cast_float_to_string.rs index deb24a3e04c..71af3074c8b 100644 --- a/native/spark-expr/benches/cast_float_to_string.rs +++ b/native/spark-expr/benches/cast_float_to_string.rs @@ -67,7 +67,7 @@ fn criterion_benchmark(c: &mut Criterion) { let size = 8192; let f64_array = create_f64_array(size); let f32_array = create_f32_array(size); - let cast_options = SparkCastOptions::new_without_timezone(EvalMode::Legacy, false); + let cast_options = SparkCastOptions::new_without_timezone(EvalMode::Legacy); let mut group = c.benchmark_group("cast_float_to_string"); group.bench_function("cast_f64_to_utf8", |b| { diff --git a/native/spark-expr/benches/cast_from_boolean.rs b/native/spark-expr/benches/cast_from_boolean.rs index caccd67e26d..48f080568aa 100644 --- a/native/spark-expr/benches/cast_from_boolean.rs +++ b/native/spark-expr/benches/cast_from_boolean.rs @@ -26,7 +26,7 @@ use std::sync::Arc; fn criterion_benchmark(c: &mut Criterion) { let expr = Arc::new(Column::new("a", 0)); let boolean_batch = create_boolean_batch(); - let spark_cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let spark_cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC"); let cast_to_i8 = Cast::new( expr.clone(), DataType::Int8, diff --git a/native/spark-expr/benches/cast_from_string.rs b/native/spark-expr/benches/cast_from_string.rs index 31cf11a84ba..cd344670127 100644 --- a/native/spark-expr/benches/cast_from_string.rs +++ b/native/spark-expr/benches/cast_from_string.rs @@ -41,7 +41,7 @@ fn criterion_benchmark(c: &mut Criterion) { let expr = Arc::new(Column::new("a", 0)); for (mode, mode_name) in EVAL_MODES { - let spark_cast_options = SparkCastOptions::new(mode, "", false); + let spark_cast_options = SparkCastOptions::new(mode, ""); let cast_to_i8 = Cast::new( expr.clone(), DataType::Int8, @@ -88,7 +88,7 @@ fn criterion_benchmark(c: &mut Criterion) { } // Benchmark decimal truncation (Legacy mode only) - let spark_cast_options = SparkCastOptions::new(EvalMode::Legacy, "", false); + let spark_cast_options = SparkCastOptions::new(EvalMode::Legacy, ""); let cast_to_i32 = Cast::new( expr.clone(), DataType::Int32, @@ -116,7 +116,7 @@ fn criterion_benchmark(c: &mut Criterion) { // str -> decimal benchmark let decimal_string_batch = create_decimal_cast_string_batch(); for (mode, mode_name) in EVAL_MODES { - let spark_cast_options = SparkCastOptions::new(mode, "", false); + let spark_cast_options = SparkCastOptions::new(mode, ""); let mut group = c.benchmark_group(format!("cast_string_to_decimal/{}", mode_name)); for (data_type, name) in [ (DataType::Decimal128(38, 10), "decimal_38_10"), @@ -143,7 +143,7 @@ fn criterion_benchmark(c: &mut Criterion) { let float_batch = create_float_string_batch(false); let float_padded_batch = create_float_string_batch(true); for (mode, mode_name) in EVAL_MODES { - let spark_cast_options = SparkCastOptions::new(mode, "", false); + let spark_cast_options = SparkCastOptions::new(mode, ""); let mut group = c.benchmark_group(format!("cast_string_to_bool_and_float/{}", mode_name)); for (data_type, name, batch) in [ (DataType::Boolean, "boolean", &bool_batch), diff --git a/native/spark-expr/benches/cast_int_to_decimal.rs b/native/spark-expr/benches/cast_int_to_decimal.rs index 8949c5c3276..a27b5d277b0 100644 --- a/native/spark-expr/benches/cast_int_to_decimal.rs +++ b/native/spark-expr/benches/cast_int_to_decimal.rs @@ -56,7 +56,7 @@ fn cast(col: &str, to: DataType, mode: EvalMode) -> Cast { Cast::new( Arc::new(Column::new(col, 0)), to, - SparkCastOptions::new_without_timezone(mode, false), + SparkCastOptions::new_without_timezone(mode), None, None, ) diff --git a/native/spark-expr/benches/cast_int_to_timestamp.rs b/native/spark-expr/benches/cast_int_to_timestamp.rs index 4479627ae66..04a62280d5a 100644 --- a/native/spark-expr/benches/cast_int_to_timestamp.rs +++ b/native/spark-expr/benches/cast_int_to_timestamp.rs @@ -27,7 +27,7 @@ const BATCH_SIZE: usize = 8192; fn criterion_benchmark(c: &mut Criterion) { // Test with UTC timezone - let spark_cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let spark_cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC"); let timestamp_type = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())); let mut group = c.benchmark_group("cast_int_to_timestamp"); diff --git a/native/spark-expr/benches/cast_nested.rs b/native/spark-expr/benches/cast_nested.rs new file mode 100644 index 00000000000..0474f2b9af2 --- /dev/null +++ b/native/spark-expr/benches/cast_nested.rs @@ -0,0 +1,141 @@ +// 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. + +//! Benchmarks for the nested `cast` paths: struct->struct (child retyped), map->map (value +//! retyped), and array->string. Each builds an 8192-row input and casts it through the public +//! `Cast` physical expression, mirroring the scalar cast benches. + +use arrow::array::{ + ArrayRef, Int32Array, ListArray, MapArray, RecordBatch, StringArray, StructArray, +}; +use arrow::buffer::OffsetBuffer; +use arrow::datatypes::{DataType, Field, Fields, Int32Type, Schema}; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::PhysicalExpr; +use datafusion_comet_spark_expr::{Cast, EvalMode, SparkCastOptions}; +use std::hint::black_box; +use std::sync::Arc; + +const NUM_ROWS: usize = 8192; +const ENTRIES_PER_ROW: usize = 2; + +/// Single-column batch named `c` wrapping `array`. +fn batch(array: ArrayRef) -> RecordBatch { + let schema = Schema::new(vec![Field::new("c", array.data_type().clone(), true)]); + RecordBatch::try_new(Arc::new(schema), vec![array]).unwrap() +} + +/// `Struct` with `NUM_ROWS` rows. +fn struct_input() -> ArrayRef { + let a = Arc::new(Int32Array::from((0..NUM_ROWS as i32).collect::>())) as ArrayRef; + let b = Arc::new(StringArray::from( + (0..NUM_ROWS).map(|i| format!("s{i}")).collect::>(), + )) as ArrayRef; + let fields = Fields::from(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + ]); + Arc::new(StructArray::new(fields, vec![a, b], None)) +} + +/// `Map` with `NUM_ROWS` rows of `ENTRIES_PER_ROW` entries each. +fn map_input() -> ArrayRef { + let total = NUM_ROWS * ENTRIES_PER_ROW; + let keys = Arc::new(StringArray::from( + (0..total).map(|i| format!("k{i}")).collect::>(), + )) as ArrayRef; + let values = Arc::new(Int32Array::from((0..total as i32).collect::>())) as ArrayRef; + let entries_fields = Fields::from(vec![ + Field::new("key_value_key", DataType::Utf8, false), + Field::new("key_value_value", DataType::Int32, true), + ]); + let entries = StructArray::new(entries_fields, vec![keys, values], None); + let offsets = OffsetBuffer::::new( + (0..=NUM_ROWS) + .map(|i| (i * ENTRIES_PER_ROW) as i32) + .collect::>() + .into(), + ); + let entries_field = Arc::new(Field::new( + "key_value", + DataType::Struct(entries.fields().clone()), + false, + )); + Arc::new(MapArray::new(entries_field, offsets, entries, None, false)) +} + +/// `List` with `NUM_ROWS` rows whose lengths cycle through 0..=3. +fn list_input() -> ArrayRef { + let data = (0..NUM_ROWS).map(|i| { + let len = i % 4; + Some( + (0..len) + .map(|j| Some((i * 4 + j) as i32)) + .collect::>(), + ) + }); + Arc::new(ListArray::from_iter_primitive::(data)) +} + +fn criterion_benchmark(c: &mut Criterion) { + let options = SparkCastOptions::new(EvalMode::Legacy, "UTC"); + let col: Arc = Arc::new(Column::new("c", 0)); + + let mut group = c.benchmark_group("cast_nested"); + + // struct -> struct: retypes the first child. + let struct_batch = batch(struct_input()); + let struct_target = DataType::Struct(Fields::from(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Utf8, true), + ])); + let cast_struct = Cast::new(col.clone(), struct_target, options.clone(), None, None); + group.bench_function("struct_to_struct", |b| { + b.iter(|| black_box(cast_struct.evaluate(black_box(&struct_batch)).unwrap())) + }); + + // map -> map: retypes the value field. + let map_batch = batch(map_input()); + let map_target = DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Int64, true), + ])), + false, + )), + false, + ); + let cast_map = Cast::new(col.clone(), map_target, options.clone(), None, None); + group.bench_function("map_to_map", |b| { + b.iter(|| black_box(cast_map.evaluate(black_box(&map_batch)).unwrap())) + }); + + // list -> Utf8: renders each list to its Spark string form. + let list_batch = batch(list_input()); + let cast_list = Cast::new(col, DataType::Utf8, options, None, None); + group.bench_function("array_to_string", |b| { + b.iter(|| black_box(cast_list.evaluate(black_box(&list_batch)).unwrap())) + }); + + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/benches/cast_non_int_numeric_timestamp.rs b/native/spark-expr/benches/cast_non_int_numeric_timestamp.rs index 18ed3041eb6..9c071aae4b6 100644 --- a/native/spark-expr/benches/cast_non_int_numeric_timestamp.rs +++ b/native/spark-expr/benches/cast_non_int_numeric_timestamp.rs @@ -26,7 +26,7 @@ use std::sync::Arc; const BATCH_SIZE: usize = 8192; fn criterion_benchmark(c: &mut Criterion) { - let spark_cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let spark_cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC"); let timestamp_type = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())); let mut group = c.benchmark_group("cast_non_int_numeric_to_timestamp"); diff --git a/native/spark-expr/benches/cast_numeric.rs b/native/spark-expr/benches/cast_numeric.rs index 5153fb7d011..41fd6006d83 100644 --- a/native/spark-expr/benches/cast_numeric.rs +++ b/native/spark-expr/benches/cast_numeric.rs @@ -28,7 +28,7 @@ const NUM_ROWS: usize = 8192; fn criterion_benchmark(c: &mut Criterion) { let batch = create_int32_batch(); let expr = Arc::new(Column::new("a", 0)); - let spark_cast_options = SparkCastOptions::new_without_timezone(EvalMode::Legacy, false); + let spark_cast_options = SparkCastOptions::new_without_timezone(EvalMode::Legacy); let cast_i32_to_i8 = Cast::new( expr.clone(), DataType::Int8, @@ -61,7 +61,7 @@ fn criterion_benchmark(c: &mut Criterion) { Cast::new( Arc::new(Column::new("a", 0)), data_type, - SparkCastOptions::new_without_timezone(EvalMode::Legacy, false), + SparkCastOptions::new_without_timezone(EvalMode::Legacy), None, None, ) diff --git a/native/spark-expr/benches/cast_string_to_date.rs b/native/spark-expr/benches/cast_string_to_date.rs index aee0fc0b03a..b2a9d7dabc5 100644 --- a/native/spark-expr/benches/cast_string_to_date.rs +++ b/native/spark-expr/benches/cast_string_to_date.rs @@ -32,7 +32,7 @@ fn criterion_benchmark(c: &mut Criterion) { let cast_to_date = Cast::new( expr, DataType::Date32, - SparkCastOptions::new(EvalMode::Legacy, "UTC", false), + SparkCastOptions::new(EvalMode::Legacy, "UTC"), None, None, ); diff --git a/native/spark-expr/benches/cast_string_to_timestamp.rs b/native/spark-expr/benches/cast_string_to_timestamp.rs index d499cec96ae..117c00592e5 100644 --- a/native/spark-expr/benches/cast_string_to_timestamp.rs +++ b/native/spark-expr/benches/cast_string_to_timestamp.rs @@ -218,7 +218,7 @@ fn criterion_benchmark(c: &mut Criterion) { let cast = Cast::new( Arc::clone(&expr), to_type.clone(), - SparkCastOptions::new(mode, timezone, false), + SparkCastOptions::new(mode, timezone), None, None, ); @@ -238,7 +238,7 @@ fn criterion_benchmark(c: &mut Criterion) { let cast = Cast::new( Arc::clone(&expr), DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), - SparkCastOptions::new_with_version(EvalMode::Legacy, "UTC", false, true), + SparkCastOptions::new_with_version(EvalMode::Legacy, "UTC", true), None, None, ); diff --git a/native/spark-expr/benches/conditional.rs b/native/spark-expr/benches/conditional.rs index 8af04c2d52f..2febcddee99 100644 --- a/native/spark-expr/benches/conditional.rs +++ b/native/spark-expr/benches/conditional.rs @@ -122,7 +122,7 @@ fn spark_cast(child: Expr, data_type: DataType) -> Expr { Arc::new(Cast::new( child, data_type, - SparkCastOptions::new(EvalMode::Legacy, "America/Los_Angeles", false), + SparkCastOptions::new(EvalMode::Legacy, "America/Los_Angeles"), None, None, )) diff --git a/native/spark-expr/benches/date_trunc.rs b/native/spark-expr/benches/date_trunc.rs index a3dd0900c61..b8342c247df 100644 --- a/native/spark-expr/benches/date_trunc.rs +++ b/native/spark-expr/benches/date_trunc.rs @@ -18,31 +18,57 @@ use arrow::array::{ArrayRef, Date32Array}; use criterion::{criterion_group, criterion_main, Criterion}; use datafusion_comet_spark_expr::date_trunc_dyn; -use std::sync::Arc; +use std::{hint::black_box, sync::Arc}; fn criterion_benchmark(c: &mut Criterion) { - let date_array = create_date_array(); + let modern_dates: ArrayRef = Arc::new(create_modern_date_array()); + let mixed_range_null_dates: ArrayRef = Arc::new(create_mixed_range_null_date_array()); let mut group = c.benchmark_group("date_trunc"); // Benchmark each truncation format for format in ["YEAR", "QUARTER", "MONTH", "WEEK"] { - let array_ref: ArrayRef = Arc::new(date_array.clone()); group.bench_function(format!("date_trunc_{}", format.to_lowercase()), |b| { - b.iter(|| date_trunc_dyn(&array_ref, format.to_string()).unwrap()); + b.iter(|| { + black_box(date_trunc_dyn(black_box(&modern_dates), format.to_string()).unwrap()) + }); }); + group.bench_function( + format!("date_trunc_{}_mixed_range_null", format.to_lowercase()), + |b| { + b.iter(|| { + black_box( + date_trunc_dyn(black_box(&mixed_range_null_dates), format.to_string()) + .unwrap(), + ) + }); + }, + ); } group.finish(); } -fn create_date_array() -> Date32Array { +fn create_modern_date_array() -> Date32Array { // Create 10000 dates spanning several years (more realistic workload) // Days since Unix epoch: range from 0 (1970-01-01) to ~19000 (2022) let dates: Vec = (0..10000).map(|i| (i * 2) % 19000).collect(); Date32Array::from(dates) } +fn create_mixed_range_null_date_array() -> Date32Array { + // Mix ordinary modern dates with dates around year 3333, which are outside + // DataFusion's TimestampNanosecond range, and nulls. + let dates: Vec> = (0..10000) + .map(|i| match i % 10 { + 0 | 1 => None, + 2 | 3 => Some(497_826 + (i * 17) % 365), + _ => Some((i * 2) % 19000), + }) + .collect(); + Date32Array::from(dates) +} + fn config() -> Criterion { Criterion::default() } diff --git a/native/spark-expr/benches/float_arrays.rs b/native/spark-expr/benches/float_arrays.rs index 44a4d78a2fc..87f0dd27e80 100644 --- a/native/spark-expr/benches/float_arrays.rs +++ b/native/spark-expr/benches/float_arrays.rs @@ -28,7 +28,8 @@ use datafusion::common::ScalarValue; use datafusion::functions_nested::remove::array_remove_all_udf; use datafusion::functions_nested::sort::array_sort_udf; use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDF}; -use datafusion_comet_spark_expr::{SparkArrayRemove, SparkSortArray}; +use datafusion_comet_spark_expr::{SparkArrayRemove, SparkFloatArrayContains, SparkSortArray}; +use datafusion_spark::function::array::array_contains::SparkArrayContains; use std::hint::black_box; use std::sync::Arc; use std::time::Duration; @@ -76,6 +77,9 @@ fn criterion_benchmark(c: &mut Criterion) { group.measurement_time(Duration::from_secs(1)); let comet_sort = ScalarUDF::new_from_impl(SparkSortArray::default()); let comet_remove = ScalarUDF::new_from_impl(SparkArrayRemove::default()); + let comet_contains = ScalarUDF::new_from_impl(SparkFloatArrayContains::default()); + // The native path other element types take: array_has plus Spark's null semantics. + let datafusion_contains = Arc::new(ScalarUDF::new_from_impl(SparkArrayContains::default())); let (datafusion_sort, datafusion_remove) = (array_sort_udf(), array_remove_all_udf()); for len in [8, 50] { for nulls in [false, true] { @@ -107,6 +111,16 @@ fn criterion_benchmark(c: &mut Criterion) { let row_values = Float64Array::from_iter_values((0..ROWS).map(|row| value(row * len + 1))); let remove_column_args = args( + vec![column.clone(), ColumnarValue::Array(Arc::new(row_values))], + &list, + ); + let contains_args = args( + vec![column.clone(), scalar(ScalarValue::Float64(Some(26.0)))], + &list, + ); + let row_values = + Float64Array::from_iter_values((0..ROWS).map(|row| value(row * len + 1))); + let contains_column_args = args( vec![column, ColumnarValue::Array(Arc::new(row_values))], &list, ); @@ -132,6 +146,20 @@ fn criterion_benchmark(c: &mut Criterion) { &remove_column_args, &remove_column_args, ), + ( + "array_contains", + &comet_contains, + &datafusion_contains, + &contains_args, + &contains_args, + ), + ( + "array_contains_column", + &comet_contains, + &datafusion_contains, + &contains_column_args, + &contains_column_args, + ), ]; for (op, comet, datafusion, comet_args, datafusion_args) in cases { let evaluate = |udf: &ScalarUDF, args: &ScalarFunctionArgs| { diff --git a/native/spark-expr/benches/float_comparison.rs b/native/spark-expr/benches/float_comparison.rs new file mode 100644 index 00000000000..aacadec6550 --- /dev/null +++ b/native/spark-expr/benches/float_comparison.rs @@ -0,0 +1,138 @@ +// 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. + +//! Compare Spark's float comparisons, which follow Spark's SQL ordering without normalizing their +//! operands, with comparing normalized copies of the operands and with DataFusion's comparison, on +//! the same data, and time `normalize_floats` on its own. All three return the same answer except +//! that DataFusion's differs on a NaN with the sign bit set, which is checked before timing; +//! special-value semantics belong in tests. + +use arrow::array::{Array, ArrayRef, Float64Array}; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::record_batch::RecordBatch; +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; +use datafusion::common::ScalarValue; +use datafusion::logical_expr::Operator; +use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion_comet_spark_expr::{ + normalize_floats, spark_comparison, FloatOperands, NormalizeNaNAndZero, +}; +use std::hint::black_box; +use std::sync::Arc; +use std::time::Duration; + +const ROWS: usize = 8192; + +/// A NaN with the sign bit set, which arithmetic produces on x86-64. +const NEGATIVE_NAN: f64 = f64::from_bits(0xfff8_0000_0000_0000); + +/// Two columns of doubles, with `special` in every tenth row of each if given. +fn batch(special: Option) -> RecordBatch { + let column = |seed: usize| -> ArrayRef { + Arc::new(Float64Array::from_iter_values((0..ROWS).map( + |i| match special { + Some(value) if i % 10 == seed => value, + _ => ((i * 7 + seed) % 1000) as f64 * 1.5, + }, + ))) + }; + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Float64, true), + Field::new("b", DataType::Float64, true), + ])); + RecordBatch::try_new(schema, vec![column(1), column(3)]).unwrap() +} + +fn criterion_benchmark(c: &mut Criterion) { + let mut group = c.benchmark_group("float_comparison"); + group.sample_size(20); + group.warm_up_time(Duration::from_millis(250)); + group.measurement_time(Duration::from_secs(1)); + for (data, special) in [ + ("no_special", None), + ("neg_zero_10pct", Some(-0.0)), + ("nan_10pct", Some(f64::NAN)), + ("sign_bit_nan_10pct", Some(NEGATIVE_NAN)), + ] { + let batch = batch(special); + let schema = batch.schema(); + let a: Arc = Arc::new(Column::new("a", 0)); + let b: Arc = Arc::new(Column::new("b", 1)); + let literal: Arc = + Arc::new(Literal::new(ScalarValue::Float64(Some(500.0)))); + for (op_name, op) in [ + ("lt", Operator::Lt), + ("eq", Operator::Eq), + ("not_distinct", Operator::IsNotDistinctFrom), + ] { + for (shape, right) in [("column_column", &b), ("column_literal", &literal)] { + let datafusion: Arc = + Arc::new(BinaryExpr::new(Arc::clone(&a), op, Arc::clone(right))); + // Comparing normalized copies of the operands, as `spark_comparison` did before it + // compared them in place. It normalized a literal while planning, and 500.0 is + // already normal. + let normalize = |operand: &Arc| -> Arc { + if operand.downcast_ref::().is_some() { + Arc::clone(operand) + } else { + NormalizeNaNAndZero::wrap_if_needed(Arc::clone(operand), &schema).unwrap() + } + }; + let normalized: Arc = + Arc::new(BinaryExpr::new(normalize(&a), op, normalize(right))); + let spark = spark_comparison( + Arc::clone(&a), + op, + Arc::clone(right), + &schema, + FloatOperands::Normalize, + ) + .unwrap(); + let evaluate = |expr: &Arc| { + expr.evaluate(&batch) + .unwrap() + .into_array(ROWS) + .unwrap() + .to_data() + }; + assert_eq!(evaluate(&spark), evaluate(&normalized)); + if special.map(f64::to_bits) != Some(NEGATIVE_NAN.to_bits()) { + assert_eq!(evaluate(&spark), evaluate(&datafusion)); + } + for (engine, expr) in [ + ("spark", &spark), + ("normalized", &normalized), + ("datafusion", &datafusion), + ] { + group.bench_function( + BenchmarkId::new(format!("{op_name}_{shape}_{engine}"), data), + |b| b.iter(|| black_box(expr.evaluate(black_box(&batch)).unwrap())), + ); + } + } + } + let column = Arc::clone(batch.column(0)); + group.bench_function(BenchmarkId::new("normalize_floats", data), |b| { + b.iter(|| black_box(normalize_floats(black_box(&column)))) + }); + } + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/benches/hll_agg.rs b/native/spark-expr/benches/hll_agg.rs new file mode 100644 index 00000000000..c933533f7c6 --- /dev/null +++ b/native/spark-expr/benches/hll_agg.rs @@ -0,0 +1,184 @@ +// 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. + +//! Benchmarks for the Comet HyperLogLog aggregates in `agg_funcs`: `hll_sketch_agg` +//! (builds a sketch from raw values) and `hll_union_agg` (merges serialized sketches). +//! Each runs as a grouped `AggregateMode::Partial` over 8192 rows x 10 batches, mirroring +//! the other aggregate benches. The union input is a column of real sketch bytes minted +//! up front with `SparkHllSketch`, since the kernel deserializes and merges each one. + +use arrow::array::{ArrayRef, BinaryBuilder, Int64Builder, RecordBatch, StringBuilder}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; +use datafusion::execution::TaskContext; +use datafusion::logical_expr::AggregateUDF; +use datafusion::physical_expr::aggregate::AggregateExprBuilder; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_comet_spark_expr::{HllSketchAgg, HllUnionAgg, SparkHllSketch}; +use futures::StreamExt; +use std::hint::black_box; +use std::sync::Arc; +use tokio::runtime::Runtime; + +/// Spark's default `lgConfigK` for `approx_count_distinct`/HLL sketches. +const LG_CONFIG_K: u8 = 12; +/// Number of distinct sketches minted for the union input; rows cycle through them. +const SKETCH_POOL: usize = 256; + +async fn agg_test( + partitions: &[Vec], + group_col: Arc, + input_exprs: Vec>, + aggregate_udf: Arc, + alias: &str, +) { + let schema = &partitions[0][0].schema(); + let scan: Arc = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(partitions, Arc::clone(schema), None).unwrap(), + ))); + let aggregate = create_aggregate(scan, group_col, input_exprs, schema, aggregate_udf, alias); + let mut stream = aggregate + .execute(0, Arc::new(TaskContext::default())) + .unwrap(); + while let Some(batch) = stream.next().await { + let _batch = batch.unwrap(); + } +} + +fn create_aggregate( + scan: Arc, + group_col: Arc, + input_exprs: Vec>, + schema: &SchemaRef, + aggregate_udf: Arc, + alias: &str, +) -> Arc { + let aggr_expr = AggregateExprBuilder::new(aggregate_udf, input_exprs) + .schema(schema.clone()) + .alias(alias) + .with_ignore_nulls(false) + .with_distinct(false) + .build() + .unwrap(); + + Arc::new( + AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(vec![(group_col, "c0".to_string())]), + vec![aggr_expr.into()], + vec![None], // no filter expressions + scan, + Arc::clone(schema), + ) + .unwrap(), + ) +} + +/// A pool of `SKETCH_POOL` distinct non-empty sketches, each covering a disjoint run of i64 +/// values so a union of several estimates to a meaningful cardinality. +fn sketch_pool() -> Vec> { + (0..SKETCH_POOL) + .map(|p| { + let mut sketch = SparkHllSketch::new(LG_CONFIG_K); + let start = (p * 100) as i64; + for v in start..start + 100 { + sketch.update_i64(v); + } + sketch.to_sketch_bytes() + }) + .collect() +} + +/// `c0` is the Utf8 group key (1024 groups); `c1` is a spread-out Int64 input for +/// `hll_sketch_agg`; `c2` is a column of serialized sketches for `hll_union_agg`. +fn create_record_batch(num_rows: usize, pool: &[Vec]) -> RecordBatch { + let mut group_builder = StringBuilder::with_capacity(num_rows, num_rows * 8); + let mut c1_builder = Int64Builder::with_capacity(num_rows); + let mut c2_builder = BinaryBuilder::new(); + for i in 0..num_rows { + group_builder.append_value(format!("group_{}", i % 1024)); + c1_builder.append_value(i as i64); + c2_builder.append_value(&pool[i % pool.len()]); + } + + let fields = vec![ + Field::new("c0", DataType::Utf8, false), + Field::new("c1", DataType::Int64, false), + Field::new("c2", DataType::Binary, false), + ]; + let columns: Vec = vec![ + Arc::new(group_builder.finish()), + Arc::new(c1_builder.finish()), + Arc::new(c2_builder.finish()), + ]; + + RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap() +} + +fn criterion_benchmark(c: &mut Criterion) { + let num_rows = 8192; + let pool = sketch_pool(); + let batch = create_record_batch(num_rows, &pool); + let batches = vec![batch; 10]; + let partitions = &[batches]; + + let c0: Arc = Arc::new(Column::new("c0", 0)); + let c1: Arc = Arc::new(Column::new("c1", 1)); + let c2: Arc = Arc::new(Column::new("c2", 2)); + + let rt = Runtime::new().unwrap(); + let mut group = c.benchmark_group("hll_agg"); + + let sketch_udf: Arc = Arc::new(AggregateUDF::new_from_impl(HllSketchAgg::new( + LG_CONFIG_K as i32, + ))); + group.bench_function("hll_sketch_agg", |b| { + b.to_async(&rt).iter(|| { + black_box(agg_test( + partitions, + c0.clone(), + vec![c1.clone()], + sketch_udf.clone(), + "hll_sketch_agg", + )) + }) + }); + + let union_udf: Arc = + Arc::new(AggregateUDF::new_from_impl(HllUnionAgg::new(false))); + group.bench_function("hll_union_agg", |b| { + b.to_async(&rt).iter(|| { + black_box(agg_test( + partitions, + c0.clone(), + vec![c2.clone()], + union_udf.clone(), + "hll_union_agg", + )) + }) + }); + + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/benches/hll_scalar.rs b/native/spark-expr/benches/hll_scalar.rs new file mode 100644 index 00000000000..d3b4c79070f --- /dev/null +++ b/native/spark-expr/benches/hll_scalar.rs @@ -0,0 +1,74 @@ +// 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. + +//! Benchmarks for the scalar HLL kernels: `hll_sketch_estimate` (Binary sketch -> Long estimate) +//! and `hll_union` (two Binary sketch columns + an `allowDifferentLgConfigK` flag -> Binary). +//! Both deserialize a sketch per row, so inputs are real sketch bytes minted up front with +//! `SparkHllSketch`. + +use arrow::array::{ArrayRef, BinaryArray, BooleanArray}; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::physical_plan::ColumnarValue; +use datafusion_comet_spark_expr::{spark_hll_sketch_estimate, spark_hll_union, SparkHllSketch}; +use std::hint::black_box; +use std::sync::Arc; + +const LG_CONFIG_K: u8 = 12; + +/// Build a Binary column of `rows` sketches, each covering a distinct run of 100 i64 values +/// offset by `base` so the two columns in a union cover different ranges. +fn sketch_column(rows: usize, base: i64) -> ArrayRef { + let sketches: Vec> = (0..rows) + .map(|i| { + let mut s = SparkHllSketch::new(LG_CONFIG_K); + let start = base + (i as i64) * 100; + for v in start..start + 100 { + s.update_i64(v); + } + s.to_sketch_bytes() + }) + .collect(); + Arc::new(BinaryArray::from_iter_values(sketches)) +} + +fn criterion_benchmark(c: &mut Criterion) { + let rows = 8192; + let left = sketch_column(rows, 0); + let right = sketch_column(rows, 1_000_000); + let allow: ArrayRef = Arc::new(BooleanArray::from(vec![false; rows])); + + let mut group = c.benchmark_group("hll_scalar"); + + group.bench_function("hll_sketch_estimate", |b| { + let args = vec![ColumnarValue::Array(Arc::clone(&left))]; + b.iter(|| black_box(spark_hll_sketch_estimate(black_box(&args)).unwrap())) + }); + + group.bench_function("hll_union", |b| { + let args = vec![ + ColumnarValue::Array(Arc::clone(&left)), + ColumnarValue::Array(Arc::clone(&right)), + ColumnarValue::Array(Arc::clone(&allow)), + ]; + b.iter(|| black_box(spark_hll_union(black_box(&args)).unwrap())) + }); + + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/benches/nested_comparison.rs b/native/spark-expr/benches/nested_comparison.rs index cf3fd44a077..3bfd0a34d38 100644 --- a/native/spark-expr/benches/nested_comparison.rs +++ b/native/spark-expr/benches/nested_comparison.rs @@ -15,8 +15,9 @@ // specific language governing permissions and limitations // under the License. -//! Compare the base DataFusion path, PR #6073's eager normalization, and Spark equality. -//! Ordinary finite inputs keep the answers identical across all three implementations. +//! Compare the base DataFusion path, eager normalization of the operands (PR #6073's for `IN`, and +//! the one `spark_comparison` used for ordering comparisons before it compared them in place), and +//! Spark's comparisons. Ordinary finite inputs keep the answers identical across all three. use arrow::array::{ArrayRef, Float64Array, ListArray}; use arrow::buffer::{NullBuffer, OffsetBuffer}; @@ -27,7 +28,9 @@ use datafusion::common::ScalarValue; use datafusion::logical_expr::Operator; use datafusion::physical_expr::expressions::{in_list, BinaryExpr, Column, Literal}; use datafusion::physical_expr::PhysicalExpr; -use datafusion_comet_spark_expr::{spark_comparison, spark_in_list, NormalizeNestedFloats}; +use datafusion_comet_spark_expr::{ + spark_comparison, spark_in_list, FloatOperands, NormalizeNestedFloats, +}; use std::hint::black_box; use std::sync::Arc; use std::time::Duration; @@ -63,11 +66,23 @@ fn expression(version: &str, mode: &str, batch: &RecordBatch) -> Arc = Arc::new(Column::new("a", 0)); let b: Arc = Arc::new(Column::new("b", 1)); - if mode == "eq" { - return if version == "new" { - spark_comparison(a, Operator::Eq, b, &schema).unwrap() - } else { - Arc::new(BinaryExpr::new(a, Operator::Eq, b)) + let normalize = |e: Arc| NormalizeNestedFloats::wrap_if_needed(e, &schema); + let comparison = match mode { + "eq" => Some(Operator::Eq), + "lt" => Some(Operator::Lt), + "not_distinct" => Some(Operator::IsNotDistinctFrom), + _ => None, + }; + if let Some(op) = comparison { + return match version { + "new" => spark_comparison(a, op, b, &schema, FloatOperands::Normalize).unwrap(), + // `=` already compared in place before ordering comparisons did. + "head" if op != Operator::Eq => Arc::new(BinaryExpr::new( + normalize(a).unwrap(), + op, + normalize(b).unwrap(), + )), + _ => Arc::new(BinaryExpr::new(a, op, b)), }; } let literal: Arc = Arc::new(Literal::new( @@ -116,7 +131,7 @@ fn benchmark(c: &mut Criterion) { Field::new("b", b.data_type().clone(), true), ])); let batch = RecordBatch::try_new(schema, vec![a, b]).unwrap(); - for mode in ["dynamic", "constant", "mixed", "eq"] { + for mode in ["dynamic", "constant", "mixed", "eq", "lt", "not_distinct"] { let expected = expression("base", mode, &batch) .evaluate(&batch) .unwrap() diff --git a/native/spark-expr/benches/timestamp_trunc.rs b/native/spark-expr/benches/timestamp_trunc.rs index 6dd5d9fad11..bc6db2de32b 100644 --- a/native/spark-expr/benches/timestamp_trunc.rs +++ b/native/spark-expr/benches/timestamp_trunc.rs @@ -15,8 +15,10 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::RecordBatch; -use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; +use arrow::array::{ + Array, DictionaryArray, Int32Array, RecordBatch, StringArray, TimestampMicrosecondArray, +}; +use arrow::datatypes::{DataType, Field, Int32Type, Schema, TimeUnit}; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; use datafusion::physical_expr::expressions::{lit, Column}; use datafusion::physical_expr::PhysicalExpr; @@ -26,37 +28,188 @@ use std::sync::Arc; #[path = "common/mod.rs"] mod common; -use common::{timestamp_micros_array, NULL_RATIOS, ROW_COUNTS}; +use common::{is_null, timestamp_micros_array, NULL_RATIOS, ROW_COUNTS}; -// Non-UTC session timezone so the benchmark exercises the `array_with_timezone` resolution path -// (the hot part of `TimestampTruncExpr`), not just the truncation kernel. -const TZ: &str = "America/Los_Angeles"; +// Cover the default UTC path and non-UTC timezone resolution with identical input shapes. const MICROS_PER_DAY: i64 = 86_400_000_000; +const BASE_MICROS: i64 = 1_704_067_200_123_456; // 2024-01-01T00:00:00.123456Z fn criterion_benchmark(c: &mut Criterion) { - let schema = Arc::new(Schema::new(vec![Field::new( - "a", - DataType::Timestamp(TimeUnit::Microsecond, Some(TZ.into())), - true, - )])); - let mut group = c.benchmark_group("timestamp_trunc"); - for format in ["YEAR", "MONTH", "DAY", "HOUR"] { - let expr = - TimestampTruncExpr::new(Arc::new(Column::new("a", 0)), lit(format), TZ.to_string()); - for rows in ROW_COUNTS { - for (null_ratio, tag) in NULL_RATIOS { - let ts = timestamp_micros_array(rows, null_ratio, Some(TZ), |i| { - (i as i64) * MICROS_PER_DAY - }); - let batch = RecordBatch::try_new(Arc::clone(&schema), vec![ts]).unwrap(); + for timezone in ["UTC", "America/Los_Angeles"] { + let schema = Arc::new(Schema::new(vec![Field::new( + "a", + DataType::Timestamp(TimeUnit::Microsecond, Some(timezone.into())), + true, + )])); + for format in [ + "YEAR", + "QUARTER", + "MONTH", + "WEEK", + "DAY", + "HOUR", + "MINUTE", + "SECOND", + "MILLISECOND", + "MICROSECOND", + ] { + let expr = TimestampTruncExpr::new( + Arc::new(Column::new("a", 0)), + lit(format), + timezone.to_string(), + true, + ); + for rows in ROW_COUNTS { + for (null_ratio, tag) in NULL_RATIOS.into_iter().chain([(0.875, "dense")]) { + // Keep every row in the modern range, including the largest batch. Vary + // sub-day values so fine-unit truncation does actual work as well. + let ts = timestamp_micros_array(rows, null_ratio, Some(timezone), |i| { + BASE_MICROS + + (i % 366) as i64 * MICROS_PER_DAY + + (i as i64 * 1_234_567).rem_euclid(MICROS_PER_DAY) + }); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![ts]).unwrap(); + group.bench_with_input( + BenchmarkId::from_parameter(format!("{timezone}/{format}/{rows}/{tag}")), + &batch, + |b, batch| b.iter(|| black_box(expr.evaluate(black_box(batch)).unwrap())), + ); + } + } + } + } + group.finish(); + + // Many repeated keys should cost only as much as truncating the distinct values. + // Keep NULL keys and unused/used overflowing entries in the same matched benchmark. + let mut group = c.benchmark_group("timestamp_trunc_dictionary"); + for rows in [8_192, 65_536] { + for (format, cardinality) in [ + ("MICROSECOND", 32), + ("SECOND", 32), + ("YEAR", 32), + ("YEAR", rows), + ] { + for (null_ratio, tag) in NULL_RATIOS.into_iter().chain([(0.875, "dense")]) { + let values = TimestampMicrosecondArray::from( + (0..cardinality) + .map(|i| BASE_MICROS + i as i64 * 1_234_567) + .collect::>(), + ) + .with_timezone("UTC"); + let keys = Int32Array::from_iter( + (0..rows) + .map(|i| (!is_null(i, null_ratio)).then_some((i % cardinality) as i32)), + ); + let input = DictionaryArray::::try_new(keys, Arc::new(values)).unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "a", + input.data_type().clone(), + true, + )])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(input)]).unwrap(); + let expr = TimestampTruncExpr::new( + Arc::new(Column::new("a", 0)), + lit(format), + "UTC".into(), + true, + ); + let cardinality_tag = if cardinality == rows { "high" } else { "low" }; group.bench_with_input( - BenchmarkId::from_parameter(format!("{format}/{rows}/{tag}")), + BenchmarkId::from_parameter(format!("{format}/{rows}/{cardinality_tag}/{tag}")), &batch, |b, batch| b.iter(|| black_box(expr.evaluate(black_box(batch)).unwrap())), ); } } + for (tag, referenced) in [("unused_overflow", false), ("used_overflow", true)] { + let values = + TimestampMicrosecondArray::from(vec![BASE_MICROS, i64::MIN]).with_timezone("UTC"); + let keys = Int32Array::from_iter((0..rows).map(|i| { + if referenced && i == rows - 1 { + 1 + } else { + 0 + } + })); + let input = DictionaryArray::::try_new(keys, Arc::new(values)).unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "a", + input.data_type().clone(), + true, + )])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(input)]).unwrap(); + let expr = TimestampTruncExpr::new( + Arc::new(Column::new("a", 0)), + lit("YEAR"), + "UTC".into(), + true, + ); + assert_eq!(expr.evaluate(&batch).is_err(), referenced); + group.bench_with_input( + BenchmarkId::from_parameter(format!("YEAR/{rows}/{tag}")), + &batch, + |b, batch| b.iter(|| black_box(expr.evaluate(black_box(batch)))), + ); + } + } + group.finish(); + + // Formats from a column: one format for every row, and every unit mixed across the rows. + // Daytime instants (16:00 to 22:00 UTC) avoid DST transitions, so older kernels can be + // measured on the same input. + let mut group = c.benchmark_group("timestamp_trunc_format_column"); + let rows = 8_192; + let units = [ + "YEAR", + "QUARTER", + "MONTH", + "WEEK", + "DAY", + "HOUR", + "MINUTE", + "SECOND", + "MILLISECOND", + "MICROSECOND", + ]; + for timezone in ["UTC", "America/Los_Angeles"] { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "a", + DataType::Timestamp(TimeUnit::Microsecond, Some(timezone.into())), + true, + ), + Field::new("fmt", DataType::Utf8, true), + ])); + let timestamps = TimestampMicrosecondArray::from_iter_values((0..rows).map(|i| { + BASE_MICROS + + (i % 366) as i64 * MICROS_PER_DAY + + 57_600_000_000 + + (i as i64 * 1_234_567) % 21_600_000_000 + })) + .with_timezone(timezone); + for (tag, unit) in [("YEAR", Some(0)), ("HOUR", Some(5)), ("mixed", None)] { + let formats = StringArray::from_iter_values( + (0..rows).map(|i| units[unit.unwrap_or(i % units.len())]), + ); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(timestamps.clone()), Arc::new(formats)], + ) + .unwrap(); + let expr = TimestampTruncExpr::new( + Arc::new(Column::new("a", 0)), + Arc::new(Column::new("fmt", 1)), + timezone.to_string(), + true, + ); + group.bench_with_input( + BenchmarkId::from_parameter(format!("{timezone}/{tag}/{rows}")), + &batch, + |b, batch| b.iter(|| black_box(expr.evaluate(black_box(batch)).unwrap())), + ); + } } group.finish(); } diff --git a/native/spark-expr/benches/to_csv.rs b/native/spark-expr/benches/to_csv.rs index 8620dd0f160..f51dde2a5d0 100644 --- a/native/spark-expr/benches/to_csv.rs +++ b/native/spark-expr/benches/to_csv.rs @@ -86,7 +86,7 @@ fn criterion_benchmark(c: &mut Criterion) { let default_null_value = ""; let default_quote = "\""; let default_escape = "\\"; - let mut cast_options = SparkCastOptions::new(EvalMode::Legacy, timezone, false); + let mut cast_options = SparkCastOptions::new(EvalMode::Legacy, timezone); cast_options.null_string = default_null_value.to_string(); let csv_write_options = CsvWriteOptions::new( default_delimiter.to_string(), diff --git a/native/spark-expr/benches/wide_decimal.rs b/native/spark-expr/benches/wide_decimal.rs index b38ddbfaa06..6de9fb5e83b 100644 --- a/native/spark-expr/benches/wide_decimal.rs +++ b/native/spark-expr/benches/wide_decimal.rs @@ -85,7 +85,7 @@ fn build_old_expr( ) -> Arc { let left_col: Arc = Arc::new(Column::new("left", 0)); let right_col: Arc = Arc::new(Column::new("right", 1)); - let cast_opts = SparkCastOptions::new_without_timezone(EvalMode::Legacy, false); + let cast_opts = SparkCastOptions::new_without_timezone(EvalMode::Legacy); let left_cast = Arc::new(Cast::new( left_col, DataType::Decimal256(p1, s1), diff --git a/native/spark-expr/src/array_funcs/array_contains.rs b/native/spark-expr/src/array_funcs/array_contains.rs new file mode 100644 index 00000000000..652e8ae24e2 --- /dev/null +++ b/native/spark-expr/src/array_funcs/array_contains.rs @@ -0,0 +1,447 @@ +// 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 std::sync::Arc; + +use super::set_bits_before; +use crate::float_semantics::{compare_floats, spark_equality}; +use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, BooleanBufferBuilder, ListArray}; +use arrow::buffer::{BooleanBuffer, NullBuffer}; +use arrow::datatypes::{ArrowPrimitiveType, DataType, Float32Type, Float64Type}; +use datafusion::common::{exec_err, utils::take_function_args, Result, ScalarValue}; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, +}; +use num::Float; + +/// Spark's `array_contains` for arrays whose elements hold floats at any depth. +/// +/// Spark tests each element against the value with `genEqual`, in which `-0.0` equals `0.0` and +/// all NaNs are equal, at any depth of an array or struct element. The result is true when an +/// element matches, null when none does and the array holds a null element, and false otherwise; +/// a null array or value gives null. datafusion-spark's `array_contains`, which Comet uses for +/// other element types, compares the bits instead, so `0.0` misses a `-0.0` and a NaN misses a +/// NaN whose bits differ. +#[derive(Debug, Hash, Eq, PartialEq)] +pub struct SparkFloatArrayContains { + signature: Signature, +} + +impl Default for SparkFloatArrayContains { + fn default() -> Self { + Self { + signature: Signature::any(2, Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for SparkFloatArrayContains { + fn name(&self) -> &str { + "spark_array_contains" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Boolean) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let [array, value] = take_function_args(self.name(), &args.args)?; + let all_scalars = matches!( + (array, value), + (ColumnarValue::Scalar(_), ColumnarValue::Scalar(_)) + ); + let rows = if all_scalars { 1 } else { args.number_rows }; + let array = array.to_array(rows)?; + // Spark arrays use Arrow's 32-bit List layout. + let Some(list) = array.as_list_opt::() else { + return exec_err!( + "spark_array_contains takes a list, got {}", + array.data_type() + ); + }; + let result = match value { + ColumnarValue::Scalar(needle) => contains_constant(list, needle)?, + ColumnarValue::Array(value) => array_contains(list, value)?, + }; + if all_scalars { + Ok(ColumnarValue::Scalar(ScalarValue::Boolean( + result.is_valid(0).then(|| result.value(0)), + ))) + } else { + Ok(ColumnarValue::Array(Arc::new(result))) + } + } +} + +fn array_contains(array: &ListArray, value: &ArrayRef) -> Result { + // Spark casts the value to the element type. Anything else fails in `spark_equality`. + match (array.value_type(), value.data_type()) { + (DataType::Float32, DataType::Float32) => Ok(contains_floats::(array, value)), + (DataType::Float64, DataType::Float64) => Ok(contains_floats::(array, value)), + _ => { + let equal = spark_equality(array.values().as_ref(), value.as_ref())?; + contains_where(array, value, equal) + } + } +} + +/// [`array_contains`] with the same value for every row, which a float array answers in one pass +/// over all the values. +fn contains_constant(array: &ListArray, needle: &ScalarValue) -> Result { + match (array.value_type(), needle) { + // A null value gives null for every row. + (_, needle) if needle.is_null() => Ok(BooleanArray::new_null(array.len())), + (DataType::Float32, ScalarValue::Float32(Some(needle))) => { + Ok(contains_float_constant::(array, *needle)) + } + (DataType::Float64, ScalarValue::Float64(Some(needle))) => { + Ok(contains_float_constant::(array, *needle)) + } + _ => array_contains(array, &needle.to_array_of_size(array.len())?), + } +} + +fn contains_float_constant(array: &ListArray, needle: T::Native) -> BooleanArray +where + T: ArrowPrimitiveType, + T::Native: Float, +{ + let offsets = array.offsets(); + let first = offsets[0] as usize; + let values = array + .values() + .slice(first, offsets[array.len()] as usize - first); + let floats = values.as_primitive::(); + let buffer = floats.values(); + // For a fixed value, Spark's genEqual is one test per element: IEEE `==`, under which the two + // zeros are equal and a NaN differs from every number, or `is_nan` for a NaN value. + let matches = if needle.is_nan() { + BooleanBuffer::collect_bool(buffer.len(), |index| buffer[index].is_nan()) + } else { + BooleanBuffer::collect_bool(buffer.len(), |index| buffer[index] == needle) + }; + // A null element never matches. + let matches = match floats.nulls() { + Some(nulls) => &matches & nulls.inner(), + None => matches, + }; + // One pass over each bitmap gives every row's match count and valid count. + let matched = set_bits_before(&matches, offsets); + let valid = floats + .nulls() + .map(|nulls| set_bits_before(nulls.inner(), offsets)); + let mut values = BooleanBufferBuilder::new(array.len()); + let mut validity = BooleanBufferBuilder::new(array.len()); + for (row, bounds) in offsets.windows(2).enumerate() { + let found = matched[row + 1] > matched[row]; + let has_null = valid + .as_ref() + .is_some_and(|valid| valid[row + 1] - valid[row] < bounds[1] - bounds[0]); + values.append(found); + validity.append(array.is_valid(row) && (found || !has_null)); + } + BooleanArray::new(values.finish(), Some(NullBuffer::new(validity.finish()))) +} + +/// Tests each row's floats against the row's value in Spark's `genEqual`, straight from the +/// flattened values. A row is null when the array or the value is. +fn contains_floats(array: &ListArray, value: &ArrayRef) -> BooleanArray +where + T: ArrowPrimitiveType, + T::Native: Float, +{ + let values = array.values().as_primitive::(); + let needles = value.as_primitive::().values(); + let (buffer, nulls) = (values.values(), values.nulls()); + let row_nulls = NullBuffer::union(array.nulls(), value.nulls()); + let mut found_rows = BooleanBufferBuilder::new(array.len()); + let mut validity = BooleanBufferBuilder::new(array.len()); + for (row, bounds) in array.offsets().windows(2).enumerate() { + let (mut found, mut has_null) = (false, false); + let row_valid = row_nulls.as_ref().is_none_or(|nulls| nulls.is_valid(row)); + if row_valid { + let needle = needles[row]; + for index in bounds[0] as usize..bounds[1] as usize { + if nulls.is_some_and(|nulls| nulls.is_null(index)) { + has_null = true; + } else if compare_floats(buffer[index], needle).is_eq() { + found = true; + break; + } + } + } + found_rows.append(found); + validity.append(row_valid && (found || !has_null)); + } + BooleanArray::new( + found_rows.finish(), + Some(NullBuffer::new(validity.finish())), + ) +} + +/// Tests each row's non-null elements with `equal(element, row)`. A row is null when the array or +/// the value is. +fn contains_where(array: &ListArray, value: &ArrayRef, equal: F) -> Result +where + F: Fn(usize, usize) -> bool, +{ + let element_nulls = array.values().logical_nulls(); + let row_nulls = NullBuffer::union(array.nulls(), value.nulls()); + let mut found_rows = BooleanBufferBuilder::new(array.len()); + let mut validity = BooleanBufferBuilder::new(array.len()); + for (row, bounds) in array.offsets().windows(2).enumerate() { + let (mut found, mut has_null) = (false, false); + let row_valid = row_nulls.as_ref().is_none_or(|nulls| nulls.is_valid(row)); + if row_valid { + for element in bounds[0] as usize..bounds[1] as usize { + if element_nulls + .as_ref() + .is_some_and(|nulls| nulls.is_null(element)) + { + has_null = true; + } else if equal(element, row) { + found = true; + break; + } + } + } + found_rows.append(found); + validity.append(row_valid && (found || !has_null)); + } + Ok(BooleanArray::new( + found_rows.finish(), + Some(NullBuffer::new(validity.finish())), + )) +} + +#[cfg(test)] +mod tests { + use super::super::test_util::list; + use super::*; + use crate::float_semantics::EDGE_VALUES; + use arrow::array::{Float64Array, Int32Array, StructArray}; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::{Field, Fields}; + use datafusion::config::ConfigOptions; + + fn invoke(array: ColumnarValue, value: ColumnarValue, rows: usize) -> Result { + let return_field = Arc::new(Field::new("result", DataType::Boolean, true)); + SparkFloatArrayContains::default().invoke_with_args(ScalarFunctionArgs { + args: vec![array, value], + arg_fields: vec![], + number_rows: rows, + return_field, + config_options: Arc::new(ConfigOptions::default()), + }) + } + + fn booleans(result: ColumnarValue) -> Vec> { + let ColumnarValue::Array(array) = result else { + panic!("expected an array"); + }; + array.as_boolean().iter().collect() + } + + /// Spark's result for one row: true if a non-null element `genEqual`s the value, else null if + /// the row holds a null, else false. + fn spark_contains(row: &[Option], value: f64) -> Option { + if row + .iter() + .any(|element| element.is_some_and(|e| compare_floats(e, value).is_eq())) + { + Some(true) + } else if row.contains(&None) { + None + } else { + Some(false) + } + } + + /// Every edge value looked up in rows of edge values, so each zero and each NaN meets every + /// other, through the constant path and the per-row path. + #[test] + fn floats_follow_spark_equality() -> Result<()> { + let rows: Vec>>> = EDGE_VALUES + .iter() + .map(|skipped| { + Some( + EDGE_VALUES + .iter() + .filter(|v| v != &skipped) + .copied() + .collect(), + ) + }) + .collect(); + let array = list(&rows); + for value in EDGE_VALUES.into_iter().flatten() { + let expected: Vec> = rows + .iter() + .map(|row| spark_contains(row.as_ref().unwrap(), value)) + .collect(); + let constant = invoke( + ColumnarValue::Array(Arc::clone(&array)), + ColumnarValue::Scalar(ScalarValue::Float64(Some(value))), + rows.len(), + )?; + assert_eq!(booleans(constant), expected, "constant {value:?}"); + let per_row = Float64Array::from(vec![Some(value); rows.len()]); + let column = invoke( + ColumnarValue::Array(Arc::clone(&array)), + ColumnarValue::Array(Arc::new(per_row)), + rows.len(), + )?; + assert_eq!(booleans(column), expected, "per row {value:?}"); + } + Ok(()) + } + + #[test] + fn null_rows_null_values_and_empty_arrays() -> Result<()> { + let array = list(&[ + None, + Some(vec![]), + Some(vec![None]), + Some(vec![Some(1.0), None]), + ]); + let constant = invoke( + ColumnarValue::Array(Arc::clone(&array)), + ColumnarValue::Scalar(ScalarValue::Float64(Some(1.0))), + 4, + )?; + assert_eq!( + booleans(constant), + vec![None, Some(false), None, Some(true)] + ); + let null_value = invoke( + ColumnarValue::Array(Arc::clone(&array)), + ColumnarValue::Scalar(ScalarValue::Float64(None)), + 4, + )?; + assert_eq!(booleans(null_value), vec![None; 4]); + let per_row = Float64Array::from(vec![Some(1.0), Some(1.0), None, Some(2.0)]); + let column = invoke( + ColumnarValue::Array(array), + ColumnarValue::Array(Arc::new(per_row)), + 4, + )?; + assert_eq!(booleans(column), vec![None, Some(false), None, None]); + Ok(()) + } + + #[test] + fn all_scalars() -> Result<()> { + let row = ScalarValue::List(Arc::new( + list(&[Some(vec![Some(-0.0), None])]) + .as_list::() + .clone(), + )); + let result = invoke( + ColumnarValue::Scalar(row.clone()), + ColumnarValue::Scalar(ScalarValue::Float64(Some(0.0))), + 1, + )?; + assert!(matches!( + result, + ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))) + )); + let result = invoke( + ColumnarValue::Scalar(row), + ColumnarValue::Scalar(ScalarValue::Float64(Some(1.0))), + 1, + )?; + assert!(matches!( + result, + ColumnarValue::Scalar(ScalarValue::Boolean(None)) + )); + Ok(()) + } + + /// Elements that are themselves arrays go through `spark_equality`, which compares their + /// floats the same way. + #[test] + fn nested_elements() -> Result<()> { + // Two rows of two arrays each: [[-0.0], [1.0]] and [[NaN, 1.0], [2.0]]. + let inner = list(&[ + Some(vec![Some(-0.0)]), + Some(vec![Some(1.0)]), + Some(vec![Some(f64::NAN), Some(1.0)]), + Some(vec![Some(2.0)]), + ]); + let array: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new("item", inner.data_type().clone(), true)), + OffsetBuffer::new(vec![0, 2, 4].into()), + inner, + None, + )); + // Looked-up values: [0.0] for the first row, [-NaN, 1.0] for the second. + let values = list(&[ + Some(vec![Some(0.0)]), + Some(vec![Some(-f64::NAN), Some(1.0)]), + ]); + let result = invoke(ColumnarValue::Array(array), ColumnarValue::Array(values), 2)?; + assert_eq!(booleans(result), vec![Some(true), Some(true)]); + Ok(()) + } + + /// Struct elements compare their float fields the same way; another field still has to match. + #[test] + fn struct_elements() -> Result<()> { + let fields = Fields::from(vec![ + Field::new("x", DataType::Float64, true), + Field::new("y", DataType::Int32, true), + ]); + let elements: ArrayRef = Arc::new(StructArray::new( + fields.clone(), + vec![ + Arc::new(Float64Array::from(vec![-0.0, f64::NAN, 2.0])), + Arc::new(Int32Array::from(vec![1, 1, 1])), + ], + None, + )); + // One row holding all three structs. + let array: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new("item", elements.data_type().clone(), true)), + OffsetBuffer::new(vec![0, 3].into()), + elements, + None, + )); + let lookup = |x: f64, y: i32| -> Result { + let value: ArrayRef = Arc::new(StructArray::new( + fields.clone(), + vec![ + Arc::new(Float64Array::from(vec![x])), + Arc::new(Int32Array::from(vec![y])), + ], + None, + )); + invoke( + ColumnarValue::Array(Arc::clone(&array)), + ColumnarValue::Array(value), + 1, + ) + }; + assert_eq!(booleans(lookup(0.0, 1)?), vec![Some(true)]); + assert_eq!(booleans(lookup(-f64::NAN, 1)?), vec![Some(true)]); + assert_eq!(booleans(lookup(0.0, 2)?), vec![Some(false)]); + Ok(()) + } +} diff --git a/native/spark-expr/src/array_funcs/array_remove.rs b/native/spark-expr/src/array_funcs/array_remove.rs index 7688435e433..9eda990c3ce 100644 --- a/native/spark-expr/src/array_funcs/array_remove.rs +++ b/native/spark-expr/src/array_funcs/array_remove.rs @@ -17,7 +17,7 @@ use std::sync::Arc; -use super::with_values; +use super::{set_bits_before, with_values}; use crate::float_semantics::{compare_floats, spark_equality}; use arrow::array::{ Array, ArrayRef, AsArray, BooleanArray, BooleanBufferBuilder, ListArray, PrimitiveArray, @@ -152,22 +152,9 @@ where } /// The offsets of the elements that `keep` keeps, for a `keep` that starts at `offsets[0]`: the -/// number of its set bits before each offset, counted a word at a time in one pass. +/// number of its set bits before each offset. fn kept_offsets(keep: &BooleanBuffer, offsets: &OffsetBuffer) -> OffsetBuffer { - let chunks = keep.inner().bit_chunks(keep.offset(), keep.len()); - let mut words = chunks.iter_padded(); - let (mut word, mut word_start, mut before_word) = (words.next().unwrap_or(0), 0, 0); - let kept = offsets.iter().map(|offset| { - let end = (offset - offsets[0]) as usize; - while end >= word_start + 64 { - before_word += word.count_ones() as usize; - word = words.next().unwrap_or(0); - word_start += 64; - } - let below_end = word & ((1u64 << (end - word_start)) - 1); - (before_word + below_end.count_ones() as usize) as i32 - }); - OffsetBuffer::new(kept.collect::>().into()) + OffsetBuffer::new(set_bits_before(keep, offsets).into()) } /// Copies the floats of each row that are null or differ from the row's value in Spark's diff --git a/native/spark-expr/src/array_funcs/mod.rs b/native/spark-expr/src/array_funcs/mod.rs index 01c19c418f9..ff0d55f0bdf 100644 --- a/native/spark-expr/src/array_funcs/mod.rs +++ b/native/spark-expr/src/array_funcs/mod.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +mod array_contains; mod array_extrema; mod array_insert; mod array_position; @@ -31,6 +32,7 @@ mod sequence; mod size; mod sort_array; +pub use array_contains::SparkFloatArrayContains; pub use array_extrema::SparkArrayExtrema; pub use array_insert::ArrayInsert; pub use array_position::SparkArrayPositionFunc; @@ -42,15 +44,37 @@ pub use flatten::SparkFlatten; pub use get_array_struct_fields::GetArrayStructFields; pub use list_extract::ListExtract; pub use list_positions::ListPositionsExpr; -pub use nested_comparison::{spark_comparison, spark_in_list}; +pub use nested_comparison::{spark_comparison, spark_in_list, FloatOperands, SparkComparison}; pub use sequence::spark_sequence; pub use size::{spark_size, SparkSizeFunc}; pub use sort_array::SparkSortArray; use arrow::array::{ArrayRef, ListArray}; -use arrow::buffer::{NullBuffer, OffsetBuffer}; +use arrow::buffer::{BooleanBuffer, NullBuffer, OffsetBuffer}; use std::sync::Arc; +/// The number of set bits of `bits` before each offset, for `bits` that starts at `offsets[0]`, +/// counted a word at a time in one pass. The difference between two consecutive entries is the +/// count within that row. +fn set_bits_before(bits: &BooleanBuffer, offsets: &OffsetBuffer) -> Vec { + let chunks = bits.inner().bit_chunks(bits.offset(), bits.len()); + let mut words = chunks.iter_padded(); + let (mut word, mut word_start, mut before_word) = (words.next().unwrap_or(0), 0, 0); + offsets + .iter() + .map(|offset| { + let end = (offset - offsets[0]) as usize; + while end >= word_start + 64 { + before_word += word.count_ones() as usize; + word = words.next().unwrap_or(0); + word_start += 64; + } + let below_end = word & ((1u64 << (end - word_start)) - 1); + (before_word + below_end.count_ones() as usize) as i32 + }) + .collect() +} + /// A list with the field of `array` and the given row nulls, holding `values` at `offsets`. fn with_values( array: &ListArray, diff --git a/native/spark-expr/src/array_funcs/nested_comparison.rs b/native/spark-expr/src/array_funcs/nested_comparison.rs index 395a5ddc3f5..585c244c0eb 100644 --- a/native/spark-expr/src/array_funcs/nested_comparison.rs +++ b/native/spark-expr/src/array_funcs/nested_comparison.rs @@ -15,29 +15,30 @@ // specific language governing permissions and limitations // under the License. -//! Spark equality for nested floating-point values without materializing normalized columns. +//! Spark comparisons for floating-point values, flat or nested in lists and structs, evaluated +//! without materializing normalized columns. use crate::float_semantics::{ - has_float_leaf, normalize_nested_floats, spark_equality, NormalizeNestedFloats, + compare_float_array_scalar, compare_float_arrays, comparison_with_nulls, + is_nested_with_float_leaf, normalize_comparison_operand, normalize_float_scalar, + normalize_floats, normalize_nested_floats, spark_comparator_ignoring_nulls, + spark_equality_ignoring_nulls, NormalizeNestedFloats, }; use arrow::array::{Array, ArrayRef, AsArray, BooleanArray}; -use arrow::buffer::{BooleanBuffer, NullBuffer}; +use arrow::buffer::{BooleanBuffer, Buffer, NullBuffer}; use arrow::compute::{not, or_kleene}; use arrow::datatypes::{DataType, Schema}; use arrow::record_batch::RecordBatch; use datafusion::common::{internal_err, DFSchema, Result, ScalarValue}; use datafusion::logical_expr::{ColumnarValue, Operator}; -use datafusion::physical_expr::expressions::{in_list, BinaryExpr, InListExpr}; +use datafusion::physical_expr::expressions::{in_list, BinaryExpr, Column, InListExpr, Literal}; use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_expr_common::datum::apply_cmp; use datafusion::physical_expr_common::physical_expr::is_volatile; use std::fmt::{Display, Formatter}; use std::hash::{Hash, Hasher}; use std::sync::Arc; -fn needs_spark_equality(dt: &DataType) -> bool { - dt.is_nested() && has_float_leaf(dt) -} - struct Operand { array: ArrayRef, scalar: bool, @@ -57,27 +58,45 @@ impl Operand { }) } + /// The SQL nulls of a comparison of `self` with `other`, neither of them a null scalar: the + /// rows where either operand is null. + fn nulls(&self, other: &Self) -> Option { + match (self.scalar, other.scalar) { + (true, true) => None, + (true, false) => other.array.nulls().cloned(), + (false, true) => self.array.nulls().cloned(), + (false, false) => NullBuffer::union(self.array.nulls(), other.array.nulls()), + } + } + + /// Whether either operand is a null scalar, which makes every row of a comparison null. + fn has_null_scalar(&self, other: &Self) -> bool { + (self.scalar && self.array.is_null(0)) || (other.scalar && other.array.is_null(0)) + } + + /// The nulls of this operand over `len` rows: a null scalar is null on every row. + fn broadcast_nulls(&self, len: usize) -> Option { + if self.scalar { + self.array.is_null(0).then(|| NullBuffer::new_null(len)) + } else { + self.array.nulls().cloned() + } + } + fn equal(&self, other: &Self, rows: usize) -> Result { let scalar = self.scalar && other.scalar; let len = if scalar { 1 } else { rows }; // Scalars use a single value; only the Boolean result is broadcast across rows. - let scalar_null = - (self.scalar && self.array.is_null(0)) || (other.scalar && other.array.is_null(0)); - let result = if scalar_null + let result = if self.has_null_scalar(other) || self.array.data_type() == &DataType::Null || other.array.data_type() == &DataType::Null { BooleanArray::new_null(len) } else { // Inner nulls take part in structural equality. Outer SQL nulls are handled here. - let equal = spark_equality(self.array.as_ref(), other.array.as_ref())?; - let nulls = match (self.scalar, other.scalar) { - (true, true) => None, - (true, false) => other.array.nulls().cloned(), - (false, true) => self.array.nulls().cloned(), - (false, false) => NullBuffer::union(self.array.nulls(), other.array.nulls()), - }; - let values = BooleanBuffer::collect_bool(len, |row| { + let equal = spark_equality_ignoring_nulls(self.array.as_ref(), other.array.as_ref())?; + let nulls = self.nulls(other); + let values = collect_valid(len, nulls.as_ref(), |row| { equal( if self.scalar { 0 } else { row }, if other.scalar { 0 } else { row }, @@ -85,17 +104,90 @@ impl Operand { }); BooleanArray::new(values, nulls) }; - if scalar { - Ok(ColumnarValue::Scalar(ScalarValue::Boolean( - if result.is_null(0) { - None + Ok(into_value(result, scalar)) + } + + /// `self op other` for an ordering operator, or for `IS [NOT] DISTINCT FROM`, in the ordering + /// of `spark_comparator`. The null-safe operators test equality with `spark_equality` instead, + /// which tells lists of different lengths apart without comparing their elements. Both compare + /// only the rows where neither side is null, and then apply the nulls of the two sides. + fn compare(&self, other: &Self, op: Operator, rows: usize) -> Result { + let scalar = self.scalar && other.scalar; + let len = if scalar { 1 } else { rows }; + let (left_scalar, right_scalar) = (self.scalar, other.scalar); + let index = move |row: usize, scalar: bool| if scalar { 0 } else { row }; + let result = match op { + Operator::IsDistinctFrom | Operator::IsNotDistinctFrom => { + let equal = + spark_equality_ignoring_nulls(self.array.as_ref(), other.array.as_ref())?; + let distinct = op == Operator::IsDistinctFrom; + let (left_nulls, right_nulls) = + (self.broadcast_nulls(len), other.broadcast_nulls(len)); + let valid = NullBuffer::union(left_nulls.as_ref(), right_nulls.as_ref()); + let values = collect_valid(len, valid.as_ref(), |row| { + equal(index(row, left_scalar), index(row, right_scalar)) != distinct + }); + comparison_with_nulls(op, values, left_nulls.as_ref(), right_nulls.as_ref()) + } + _ => { + if !matches!( + op, + Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq + ) { + return internal_err!("Unsupported operator for a nested comparison: {op}"); + } + if self.has_null_scalar(other) { + BooleanArray::new_null(len) } else { - Some(result.value(0)) - }, - ))) + let compare = + spark_comparator_ignoring_nulls(self.array.as_ref(), other.array.as_ref())?; + let compare_row = + |row: usize| compare(index(row, left_scalar), index(row, right_scalar)); + let nulls = self.nulls(other); + let valid = nulls.as_ref(); + // One loop per operator, so that the test on the ordering is inlined. + let values = match op { + Operator::Lt => collect_valid(len, valid, |row| compare_row(row).is_lt()), + Operator::LtEq => collect_valid(len, valid, |row| compare_row(row).is_le()), + Operator::Gt => collect_valid(len, valid, |row| compare_row(row).is_gt()), + _ => collect_valid(len, valid, |row| compare_row(row).is_ge()), + }; + BooleanArray::new(values, nulls) + } + } + }; + Ok(into_value(result, scalar)) + } +} + +/// `test(row)` for each of `len` rows, as a bitmap. A row that `valid` marks null is not tested +/// and comes out unset, which skips comparing the values under a null list or struct without +/// testing the null buffers on every row. +fn collect_valid( + len: usize, + valid: Option<&NullBuffer>, + test: impl Fn(usize) -> bool, +) -> BooleanBuffer { + let Some(valid) = valid.filter(|valid| valid.null_count() > 0) else { + return BooleanBuffer::collect_bool(len, test); + }; + let mut words = vec![0u64; len.div_ceil(64)]; + for row in valid.valid_indices() { + words[row / 64] |= (test(row) as u64) << (row % 64); + } + BooleanBuffer::new(Buffer::from_vec(words), 0, len) +} + +/// `result` as an array, or as a scalar when both operands were scalars. +fn into_value(result: BooleanArray, scalar: bool) -> ColumnarValue { + if scalar { + ColumnarValue::Scalar(ScalarValue::Boolean(if result.is_null(0) { + None } else { - Ok(ColumnarValue::Array(Arc::new(result))) - } + Some(result.value(0)) + })) + } else { + ColumnarValue::Array(Arc::new(result)) } } @@ -219,33 +311,325 @@ impl PhysicalExpr for NestedPredicate { } } -/// Build equality after the planner has reconciled nested operand nullability. +/// A comparison of two Float32 or Float64 operands, or of two lists or structs with a float leaf, +/// in Spark's SQL ordering: `-0.0` equals `0.0`, all NaNs are equal, and NaN sorts above every +/// other value. It reads the operands as they are, without normalized copies of them. +#[derive(Debug, Eq)] +pub struct SparkComparison { + left: Arc, + op: Operator, + right: Arc, + /// Whether the operands are lists or structs rather than Float32 or Float64 values. + nested: bool, +} + +impl SparkComparison { + pub fn left(&self) -> &Arc { + &self.left + } + + pub fn op(&self) -> Operator { + self.op + } + + pub fn right(&self) -> &Arc { + &self.right + } + + /// Whether the comparison cannot fail for any value: when both operands have the same type. + /// The float kernels and the nested comparators fail only on operands of different types, + /// such as a dictionary-encoded leaf against a plain one, which the logical type check in + /// [`spark_comparison`] lets through for lists and structs. + pub(crate) fn is_infallible(&self, input_schema: &Schema) -> bool { + match ( + self.left.data_type(input_schema), + self.right.data_type(input_schema), + ) { + (Ok(left), Ok(right)) => left == right, + _ => false, + } + } +} + +impl PartialEq for SparkComparison { + fn eq(&self, other: &Self) -> bool { + self.left.eq(&other.left) + && self.op == other.op + && self.right.eq(&other.right) + && self.nested == other.nested + } +} + +impl Hash for SparkComparison { + fn hash(&self, state: &mut H) { + self.left.hash(state); + self.op.hash(state); + self.right.hash(state); + self.nested.hash(state); + } +} + +impl Display for SparkComparison { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "SparkComparison({} {} {})", + self.left, self.op, self.right + ) + } +} + +impl PhysicalExpr for SparkComparison { + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + Display::fmt(self, f) + } + + fn data_type(&self, _: &Schema) -> Result { + Ok(DataType::Boolean) + } + + fn nullable(&self, schema: &Schema) -> Result { + if matches!( + self.op, + Operator::IsDistinctFrom | Operator::IsNotDistinctFrom + ) { + return Ok(false); + } + Ok(self.left.nullable(schema)? || self.right.nullable(schema)?) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + let left = self.left.evaluate(batch)?; + let right = self.right.evaluate(batch)?; + if self.nested { + Operand::new(left)?.compare(&Operand::new(right)?, self.op, batch.num_rows()) + } else { + compare_flat(self.op, left, right) + } + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.left, &self.right] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + let Ok([left, right]) = <[_; 2]>::try_from(children) else { + return internal_err!("SparkComparison expects two children"); + }; + Ok(Arc::new(Self { + left, + op: self.op, + right, + nested: self.nested, + })) + } +} + +/// `left op right` for Float32 or Float64 operands. An operand in any other layout, which these +/// types do not produce, is normalized and compared with Arrow's kernels instead. +fn compare_flat(op: Operator, left: ColumnarValue, right: ColumnarValue) -> Result { + let result = match (&left, &right) { + (ColumnarValue::Array(l), ColumnarValue::Array(r)) => { + compare_float_arrays(op, l.as_ref(), r.as_ref())? + } + (ColumnarValue::Array(l), ColumnarValue::Scalar(r)) => { + compare_float_array_scalar(op, l.as_ref(), r)? + } + (ColumnarValue::Scalar(l), ColumnarValue::Array(r)) => match op.swap() { + Some(swapped) => compare_float_array_scalar(swapped, r.as_ref(), l)?, + None => None, + }, + (ColumnarValue::Scalar(l), ColumnarValue::Scalar(r)) => { + if let Some(result) = + compare_float_arrays(op, l.to_array()?.as_ref(), r.to_array()?.as_ref())? + { + return Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( + &result, 0, + )?)); + } + None + } + }; + match result { + Some(result) => Ok(ColumnarValue::Array(Arc::new(result))), + None => apply_cmp(op, &normalize_value(left), &normalize_value(right)), + } +} + +fn normalize_value(value: ColumnarValue) -> ColumnarValue { + match value { + ColumnarValue::Array(array) => ColumnarValue::Array(normalize_floats(&array)), + ColumnarValue::Scalar(value) => ColumnarValue::Scalar(normalize_float_scalar(value)), + } +} + +/// How [`spark_comparison`] treats floating-point operands. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum FloatOperands { + /// Compare them in Spark's SQL ordering. A literal is normalized while planning, and other + /// operands are compared as they are. + Normalize, + /// Leave a Float32 or Float64 column compared with a literal other than NaN to a plain + /// [`BinaryExpr`] of the two as they are, and compare every other shape in Spark's SQL + /// ordering. Only a scan's pushed-down data filters use this, and only when the Parquet reader + /// prunes with them but does not filter rows: Parquet pruning recognizes a column compared with + /// a literal but not a Spark comparison, and Spark's Filter above the scan applies Spark's + /// semantics to every row. A NaN literal gets a Spark comparison too, giving up pruning, + /// because a stored NaN can have other bits than the literal and a writer can leave NaNs out + /// of the column statistics. With row-level pushdown the reader would drop the rows a plain + /// comparison rejects, including a stored NaN with the sign bit set, which Arrow orders below + /// every other value, so the data filters use [`FloatOperands::Normalize`] there instead. + Raw, +} + +/// Builds a comparison with Spark's SQL ordering for floats, in which `-0.0` equals `0.0`, all +/// NaNs are equal and NaN sorts above every other value, at any depth of a list or struct. +/// +/// Arrow compares floats by IEEE 754 total order instead, so a [`SparkComparison`] compares +/// Float32 and Float64 operands, and lists and structs with a float leaf, itself, without +/// normalized copies of them: flat floats with branch-free kernels over the value buffers, and +/// nested values with `spark_comparator`. Nested `=` and `<>` compare with `spark_equality`, which +/// also serves `IN`. A literal is normalized while planning. Operands of any other type, such as +/// dictionary-encoded floats, go through [`normalize_comparison_operand`] and a plain +/// [`BinaryExpr`], and so does any other operator, such as `AND`. +/// +/// The planner reconciles the nullability of nested operands before calling this. pub fn spark_comparison( left: Arc, op: Operator, right: Arc, schema: &Schema, + float_operands: FloatOperands, ) -> Result> { + use Operator::*; + if !matches!( + op, + Eq | NotEq | Lt | LtEq | Gt | GtEq | IsDistinctFrom | IsNotDistinctFrom + ) { + return Ok(Arc::new(BinaryExpr::new(left, op, right))); + } // An operand whose type does not resolve against this schema falls back to the plain // comparison, the way `reconcile_nested_comparison_types` already leaves such operands alone. - let nested = matches!(op, Operator::Eq | Operator::NotEq) - && match (left.data_type(schema), right.data_type(schema)) { - (Ok(lt), Ok(_)) => needs_spark_equality(<), - _ => false, - }; - if nested { - validate_types(&left, std::slice::from_ref(&right), schema)?; - Ok(Arc::new(NestedPredicate { - value: left, - candidates: vec![right], - negated: op == Operator::NotEq, - membership: false, - })) + let (Ok(left_type), Ok(right_type)) = (left.data_type(schema), right.data_type(schema)) else { + return Ok(Arc::new(BinaryExpr::new(left, op, right))); + }; + if is_nested_with_float_leaf(&left_type) { + if matches!(op, Eq | NotEq) { + validate_types(&left, std::slice::from_ref(&right), schema)?; + return Ok(Arc::new(NestedPredicate { + value: left, + candidates: vec![right], + negated: op == NotEq, + membership: false, + })); + } + if DFSchema::datatype_is_logically_equal(&left_type, &right_type) { + return Ok(Arc::new(SparkComparison { + left: normalize_literal(left, schema)?, + op, + right: normalize_literal(right, schema)?, + nested: true, + })); + } + } + if float_operands == FloatOperands::Raw { + if let Some(comparison) = raw_float_comparison(&left, op, &right, schema) { + return Ok(comparison); + } + } + if matches!(left_type, DataType::Float32 | DataType::Float64) && left_type == right_type { + return Ok(Arc::new(SparkComparison { + left: normalize_literal(left, schema)?, + op, + right: normalize_literal(right, schema)?, + nested: false, + })); + } + Ok(Arc::new(BinaryExpr::new( + normalize_comparison_operand(left, schema)?, + op, + normalize_comparison_operand(right, schema)?, + ))) +} + +/// Normalizes `expr` while planning if it is a literal, which keeps it a literal. Any other +/// operand is returned as is. +fn normalize_literal( + expr: Arc, + schema: &Schema, +) -> Result> { + if expr.downcast_ref::().is_some() { + normalize_comparison_operand(expr, schema) } else { - Ok(Arc::new(BinaryExpr::new(left, op, right))) + Ok(expr) } } +/// The comparison that [`FloatOperands::Raw`] builds for a Float32 or Float64 column compared with +/// a literal other than NaN: the column and the literal as they are. `None` for any other operands. +/// +/// Statistics pruning compares `-0.0` and `0.0` as equal, but a bloom filter probe hashes the +/// literal's bits, so `=` against a zero literal becomes `=` against either zero. +fn raw_float_comparison( + left: &Arc, + op: Operator, + right: &Arc, + schema: &Schema, +) -> Option> { + let (literal, literal_on_left) = match ( + left.downcast_ref::(), + right.downcast_ref::(), + ) { + (None, Some(literal)) if is_float_column(left, schema) => (literal, false), + (Some(literal), None) if is_float_column(right, schema) => (literal, true), + _ => return None, + }; + let (negative_zero, positive_zero) = match literal.value() { + ScalarValue::Float32(Some(v)) if v.is_nan() => return None, + ScalarValue::Float64(Some(v)) if v.is_nan() => return None, + ScalarValue::Float32(Some(v)) if op == Operator::Eq && *v == 0.0 => ( + ScalarValue::Float32(Some(-0.0)), + ScalarValue::Float32(Some(0.0)), + ), + ScalarValue::Float64(Some(v)) if op == Operator::Eq && *v == 0.0 => ( + ScalarValue::Float64(Some(-0.0)), + ScalarValue::Float64(Some(0.0)), + ), + _ => { + return Some(Arc::new(BinaryExpr::new( + Arc::clone(left), + op, + Arc::clone(right), + ))) + } + }; + let equal_to = |zero: ScalarValue| -> Arc { + let zero: Arc = Arc::new(Literal::new(zero)); + if literal_on_left { + Arc::new(BinaryExpr::new(zero, op, Arc::clone(right))) + } else { + Arc::new(BinaryExpr::new(Arc::clone(left), op, zero)) + } + }; + Some(Arc::new(BinaryExpr::new( + equal_to(negative_zero), + Operator::Or, + equal_to(positive_zero), + ))) +} + +/// Whether `expr` is a Float32 or Float64 column, the operand that Parquet pruning reads. +fn is_float_column(expr: &Arc, schema: &Schema) -> bool { + expr.downcast_ref::().is_some() + && matches!( + expr.data_type(schema), + Ok(DataType::Float32 | DataType::Float64) + ) +} + fn validate_types( value: &Arc, candidates: &[Arc], @@ -268,7 +652,7 @@ pub fn spark_in_list( negated: bool, schema: &Schema, ) -> Result> { - if !needs_spark_equality(&value.data_type(schema)?) || candidates.is_empty() { + if !is_nested_with_float_leaf(&value.data_type(schema)?) || candidates.is_empty() { return in_list(value, candidates, &negated, schema); } validate_types(&value, &candidates, schema)?; @@ -441,6 +825,7 @@ mod tests { }, Arc::clone(&b), batch.schema().as_ref(), + FloatOperands::Normalize, )?; let inside = spark_in_list( Arc::clone(&a), @@ -458,6 +843,7 @@ mod tests { Operator::Eq, Arc::clone(&lit), batch.schema().as_ref(), + FloatOperands::Normalize, )?; let inside = spark_in_list( Arc::clone(&a), @@ -471,6 +857,7 @@ mod tests { Operator::Eq, Arc::clone(&a), batch.schema().as_ref(), + FloatOperands::Normalize, )?; assert_eq!(results(&eq, &batch), results(&reversed, &batch)); } @@ -558,6 +945,7 @@ mod tests { Operator::Eq, Arc::clone(&zero), batch.schema().as_ref(), + FloatOperands::Normalize, )?; assert!(matches!( scalar.evaluate(&RecordBatch::new_empty(batch.schema()))?, @@ -568,6 +956,7 @@ mod tests { Operator::Eq, Arc::clone(&b), batch.schema().as_ref(), + FloatOperands::Normalize, )?; assert!(eq.nullable(batch.schema().as_ref())?); assert_eq!(eq.data_type(batch.schema().as_ref())?, DataType::Boolean); @@ -623,7 +1012,8 @@ mod tests { Arc::clone(&primitive), Operator::Eq, primitive, - batch.schema().as_ref() + batch.schema().as_ref(), + FloatOperands::Normalize, )? .as_ref() .is::()); @@ -635,11 +1025,15 @@ mod tests { (Arc::clone(&unresolved), Arc::clone(&a)), (Arc::clone(&a), Arc::clone(&unresolved)), ] { - assert!( - spark_comparison(l, Operator::Eq, r, batch.schema().as_ref())? - .as_ref() - .is::() - ); + assert!(spark_comparison( + l, + Operator::Eq, + r, + batch.schema().as_ref(), + FloatOperands::Normalize + )? + .as_ref() + .is::()); } Ok(()) } @@ -762,4 +1156,607 @@ mod tests { assert!(expr.evaluate(&batch).is_err()); Ok(()) } + + const COMPARISONS: [Operator; 8] = [ + Operator::Eq, + Operator::NotEq, + Operator::Lt, + Operator::LtEq, + Operator::Gt, + Operator::GtEq, + Operator::IsDistinctFrom, + Operator::IsNotDistinctFrom, + ]; + + /// Spark's answer for `left op right`, given the ordering of two non-null values. + fn spark_answer(op: Operator, ordering: Option>) -> Option { + use Operator::*; + match ordering { + // Both non-null. + Some(Some(ord)) => Some(match op { + Eq | IsNotDistinctFrom => ord.is_eq(), + NotEq | IsDistinctFrom => ord.is_ne(), + Lt => ord.is_lt(), + LtEq => ord.is_le(), + Gt => ord.is_gt(), + GtEq => ord.is_ge(), + _ => unreachable!(), + }), + // Both null. + Some(None) => match op { + IsNotDistinctFrom => Some(true), + IsDistinctFrom => Some(false), + _ => None, + }, + // One null. + None => match op { + IsNotDistinctFrom => Some(false), + IsDistinctFrom => Some(true), + _ => None, + }, + } + } + + fn float_ordering(l: Option, r: Option) -> Option> { + match (l, r) { + (Some(l), Some(r)) => Some(Some(crate::float_semantics::compare_floats(l, r))), + (None, None) => Some(None), + _ => None, + } + } + + /// Every pair of edge values, under every operator, as two columns and as a column and a + /// literal on either side, for Float64 and Float32. + #[test] + fn float_operands_follow_spark_ordering() -> Result<()> { + let values = [ + Some(f64::NEG_INFINITY), + Some(-1.0), + Some(-0.0), + Some(0.0), + Some(f64::MIN_POSITIVE), + Some(f64::INFINITY), + Some(f64::NAN), + // A NaN with the sign bit set, as arithmetic produces on x86-64. + Some(f64::from_bits(0xfff8_0000_0000_0000)), + // A NaN with a payload. + Some(f64::from_bits(0x7ff0_0000_0000_0001)), + None, + ]; + // The Float32 pass computes Spark's answers from the values as Float32 holds them. + let to_f32 = |v: Option| v.map(|v| v as f32 as f64); + for float32 in [false, true] { + let values: Vec<_> = if float32 { + values.iter().map(|&v| to_f32(v)).collect() + } else { + values.to_vec() + }; + let pairs: Vec<_> = values + .iter() + .flat_map(|&l| values.iter().map(move |&r| (l, r))) + .collect(); + let left: Vec<_> = pairs.iter().map(|p| p.0).collect(); + let right: Vec<_> = pairs.iter().map(|p| p.1).collect(); + let array = |v: &[Option]| -> ArrayRef { + if float32 { + Arc::new(arrow::array::Float32Array::from( + v.iter().map(|v| v.map(|v| v as f32)).collect::>(), + )) + } else { + Arc::new(Float64Array::from(v.to_vec())) + } + }; + let batch = batch(array(&left), array(&right)); + let schema = batch.schema(); + let (a, b) = columns(); + for op in COMPARISONS { + let expr = spark_comparison( + Arc::clone(&a), + op, + Arc::clone(&b), + &schema, + FloatOperands::Normalize, + )?; + let expected: Vec<_> = pairs + .iter() + .map(|&(l, r)| spark_answer(op, float_ordering(l, r))) + .collect(); + assert_eq!(results(&expr, &batch), expected, "a {op} b"); + } + // The first rows pair `values[0]` with each value in turn, so row `i` of column `b` + // holds `values[i]`. + for (i, &value) in values.iter().enumerate() { + let lit = literal(batch.column(1), i); + for op in COMPARISONS { + let column_first = spark_comparison( + Arc::clone(&a), + op, + Arc::clone(&lit), + &schema, + FloatOperands::Normalize, + )?; + let expected: Vec<_> = left + .iter() + .map(|&l| spark_answer(op, float_ordering(l, value))) + .collect(); + assert_eq!(results(&column_first, &batch), expected, "a {op} {lit}"); + let literal_first = spark_comparison( + Arc::clone(&lit), + op, + Arc::clone(&a), + &schema, + FloatOperands::Normalize, + )?; + let expected: Vec<_> = left + .iter() + .map(|&l| spark_answer(op, float_ordering(value, l))) + .collect(); + assert_eq!(results(&literal_first, &batch), expected, "{lit} {op} a"); + } + } + } + Ok(()) + } + + /// A literal operand is normalized while planning, so the comparison stays `column op + /// literal`, and a column is compared as it is. With `FloatOperands::Raw`, as for a scan's + /// data filters, a float column compared with a literal stays a plain `BinaryExpr` of the two + /// as they are. + #[test] + fn float_operand_shapes() -> Result<()> { + let batch = batch( + Arc::new(Float64Array::from(vec![1.0])), + Arc::new(Float64Array::from(vec![1.0])), + ); + let schema = batch.schema(); + let negative_zero: Arc = + Arc::new(Literal::new(ScalarValue::Float64(Some(-0.0)))); + let expr = spark_comparison( + Arc::new(Column::new("a", 0)), + Operator::Lt, + Arc::clone(&negative_zero), + &schema, + FloatOperands::Normalize, + )?; + let comparison = expr.downcast_ref::().unwrap(); + assert!(comparison.left().downcast_ref::().is_some()); + let folded = comparison.right().downcast_ref::().unwrap(); + assert!(matches!(folded.value(), ScalarValue::Float64(Some(v)) if v.to_bits() == 0)); + + // With `FloatOperands::Raw`, a float column compared with a literal keeps both as they + // are, on either side. + let expr = spark_comparison( + Arc::new(Column::new("a", 0)), + Operator::Lt, + Arc::clone(&negative_zero), + &schema, + FloatOperands::Raw, + )?; + let binary = expr.downcast_ref::().unwrap(); + assert!(binary.left().downcast_ref::().is_some()); + assert!(Arc::ptr_eq(binary.right(), &negative_zero)); + let expr = spark_comparison( + Arc::clone(&negative_zero), + Operator::Lt, + Arc::new(Column::new("a", 0)), + &schema, + FloatOperands::Raw, + )?; + let binary = expr.downcast_ref::().unwrap(); + assert!(Arc::ptr_eq(binary.left(), &negative_zero)); + assert!(binary.right().downcast_ref::().is_some()); + + // Any other shape follows Spark's ordering even with `FloatOperands::Raw`, because a + // reader with row-level pushdown drops the rows it rejects. + let (a, b) = columns(); + let expr = spark_comparison( + Arc::clone(&a), + Operator::Lt, + Arc::clone(&b), + &schema, + FloatOperands::Raw, + )?; + let comparison = expr.downcast_ref::().unwrap(); + assert!(Arc::ptr_eq(comparison.left(), &a)); + assert!(Arc::ptr_eq(comparison.right(), &b)); + assert_eq!(comparison.op(), Operator::Lt); + + // Other operators are not comparisons and keep their operands. + let expr = spark_comparison( + Arc::clone(&a), + Operator::Plus, + Arc::clone(&b), + &schema, + FloatOperands::Normalize, + )?; + let binary = expr.downcast_ref::().unwrap(); + assert!(Arc::ptr_eq(binary.left(), &a)); + assert!(Arc::ptr_eq(binary.right(), &b)); + + // Dictionary-encoded floats are normalized and compared by Arrow, as before. + let dictionary: ArrayRef = Arc::new(arrow::array::DictionaryArray::new( + Int32Array::from(vec![0, 1]), + Arc::new(Float64Array::from(vec![-0.0, 1.0])), + )); + let batch = self::batch(Arc::clone(&dictionary), dictionary); + let expr = spark_comparison( + Arc::clone(&a), + Operator::Lt, + Arc::clone(&b), + batch.schema().as_ref(), + FloatOperands::Normalize, + )?; + assert!(expr.downcast_ref::().is_some()); + assert_eq!(results(&expr, &batch), vec![Some(false), Some(false)]); + Ok(()) + } + + /// A bloom filter holds the bits of each value, so with `FloatOperands::Raw` an `=` against + /// either zero probes the filter for both zeros, while other operators keep the literal. A + /// NaN literal gets a Spark comparison instead, because no single literal stands for every + /// NaN a file can hold. + #[test] + fn raw_float_comparison_literals() -> Result<()> { + use datafusion::physical_expr::utils::{Guarantee, LiteralGuarantee}; + use std::collections::HashSet; + let schema = Schema::new(vec![ + Field::new("f", DataType::Float32, true), + Field::new("d", DataType::Float64, true), + ]); + let scalar = |value: f64, data_type: &DataType| match data_type { + DataType::Float32 => ScalarValue::Float32(Some(value as f32)), + _ => ScalarValue::Float64(Some(value)), + }; + for (index, name, data_type) in [(0, "f", DataType::Float32), (1, "d", DataType::Float64)] { + let column: Arc = Arc::new(Column::new(name, index)); + for zero in [-0.0, 0.0] { + let literal: Arc = + Arc::new(Literal::new(scalar(zero, &data_type))); + for (left, right) in [(&column, &literal), (&literal, &column)] { + let expr = spark_comparison( + Arc::clone(left), + Operator::Eq, + Arc::clone(right), + &schema, + FloatOperands::Raw, + )?; + let guarantees = LiteralGuarantee::analyze(&expr); + assert_eq!(guarantees.len(), 1, "{expr}"); + assert_eq!(guarantees[0].guarantee, Guarantee::In, "{expr}"); + assert_eq!(guarantees[0].column.name(), name, "{expr}"); + assert_eq!( + guarantees[0].literals, + HashSet::from([scalar(-0.0, &data_type), scalar(0.0, &data_type)]), + "{expr}" + ); + } + let expr = spark_comparison( + Arc::clone(&column), + Operator::GtEq, + Arc::clone(&literal), + &schema, + FloatOperands::Raw, + )?; + let binary = expr.downcast_ref::().unwrap(); + assert!(Arc::ptr_eq(binary.left(), &column)); + assert!(Arc::ptr_eq(binary.right(), &literal)); + } + } + + let column: Arc = Arc::new(Column::new("d", 1)); + for nan in [f64::NAN, f64::from_bits(0xfff8_0000_0000_0000)] { + let expr = spark_comparison( + Arc::clone(&column), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Float64(Some(nan)))), + &schema, + FloatOperands::Raw, + )?; + let comparison = expr.downcast_ref::().unwrap(); + assert!(Arc::ptr_eq(comparison.left(), &column)); + let folded = comparison.right().downcast_ref::().unwrap(); + assert!( + matches!(folded.value(), ScalarValue::Float64(Some(v)) if v.to_bits() == f64::NAN.to_bits()) + ); + } + Ok(()) + } + + /// Columns that start at an offset, and scalar operands that only appear when the batch is + /// evaluated, so that planning cannot normalize them. + #[test] + fn float_operands_at_an_offset_and_runtime_scalars() -> Result<()> { + use crate::float_semantics::EDGE_VALUES; + let pairs: Vec<_> = EDGE_VALUES + .iter() + .flat_map(|&l| EDGE_VALUES.iter().map(move |&r| (l, r))) + .collect(); + // Pad both sides with a row in front, then slice it off. + let left: Vec<_> = std::iter::once(Some(5.0)) + .chain(pairs.iter().map(|p| p.0)) + .collect(); + let right: Vec<_> = std::iter::once(None) + .chain(pairs.iter().map(|p| p.1)) + .collect(); + for float32 in [false, true] { + let array = |values: &[Option]| -> ArrayRef { + let array: ArrayRef = if float32 { + Arc::new(arrow::array::Float32Array::from( + values + .iter() + .map(|v| v.map(|v| v as f32)) + .collect::>(), + )) + } else { + Arc::new(Float64Array::from(values.to_vec())) + }; + array.slice(1, values.len() - 1) + }; + let batch = batch(array(&left), array(&right)); + let schema = batch.schema(); + let (a, b) = columns(); + // A scalar that evaluation produces as it is, unlike a literal. + let runtime_scalar = |value: Option| -> Arc { + let value = if float32 { + ScalarValue::Float32(value.map(|v| v as f32)) + } else { + ScalarValue::Float64(value) + }; + Arc::new(Probe { + child: Arc::new(Literal::new(value)), + volatile: false, + calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + }) + }; + for op in COMPARISONS { + let expr = spark_comparison( + Arc::clone(&a), + op, + Arc::clone(&b), + &schema, + FloatOperands::Normalize, + )?; + let expected: Vec<_> = pairs + .iter() + .map(|&(l, r)| spark_answer(op, float_ordering(l, r))) + .collect(); + assert_eq!(results(&expr, &batch), expected, "a {op} b, f32={float32}"); + for &value in &EDGE_VALUES { + let expr = spark_comparison( + Arc::clone(&a), + op, + runtime_scalar(value), + &schema, + FloatOperands::Normalize, + )?; + let expected: Vec<_> = pairs + .iter() + .map(|&(l, _)| spark_answer(op, float_ordering(l, value))) + .collect(); + assert_eq!(results(&expr, &batch), expected, "a {op} {value:?}"); + let expr = spark_comparison( + runtime_scalar(value), + op, + Arc::clone(&a), + &schema, + FloatOperands::Normalize, + )?; + let expected: Vec<_> = pairs + .iter() + .map(|&(l, _)| spark_answer(op, float_ordering(value, l))) + .collect(); + assert_eq!(results(&expr, &batch), expected, "{value:?} {op} a"); + // Two scalars give a scalar. + for &other in &EDGE_VALUES { + let expr = spark_comparison( + runtime_scalar(value), + op, + runtime_scalar(other), + &schema, + FloatOperands::Normalize, + )?; + let ColumnarValue::Scalar(ScalarValue::Boolean(actual)) = + expr.evaluate(&batch)? + else { + panic!("{value:?} {op} {other:?} must be a Boolean scalar"); + }; + assert_eq!( + actual, + spark_answer(op, float_ordering(value, other)), + "{value:?} {op} {other:?}" + ); + } + } + } + } + Ok(()) + } + + /// An operand that evaluates to another layout than its Float32 or Float64 type, which the + /// planner does not produce, is normalized and compared by Arrow instead of failing. + #[test] + fn unexpected_layouts_fall_back_to_normalizing() -> Result<()> { + let dictionary: ArrayRef = Arc::new(arrow::array::DictionaryArray::new( + Int32Array::from(vec![0, 1, 0]), + Arc::new(Float64Array::from(vec![1.0, 2.0])), + )); + let result = compare_flat( + Operator::Lt, + ColumnarValue::Array(dictionary), + ColumnarValue::Scalar(ScalarValue::Float64(Some(1.5))), + )? + .into_array(3)?; + assert_eq!( + result.as_boolean().iter().collect::>(), + vec![Some(true), Some(false), Some(true)] + ); + Ok(()) + } + + /// A list or struct literal on either side, a null one, and two literals, against the + /// orderings `spark_comparator` gives. + #[test] + fn nested_float_operands_with_literals() -> Result<()> { + use crate::float_semantics::spark_comparator; + let lists: ArrayRef = Arc::new(ListArray::from_iter_primitive::([ + Some(vec![Some(-0.0)]), + Some(vec![Some(0.0), Some(1.0)]), + Some(vec![Some(f64::from_bits(0xfff8_0000_0000_0000))]), + Some(vec![Some(f64::INFINITY)]), + Some(vec![None]), + Some(vec![]), + None, + ])); + let batch = batch(Arc::clone(&lists), Arc::clone(&lists)); + let schema = batch.schema(); + let (a, _) = columns(); + let null: Arc = + Arc::new(Literal::new(ScalarValue::try_from(lists.data_type())?)); + for op in COMPARISONS { + for row in 0..lists.len() { + let lit = literal(&lists, row); + let value = lit.evaluate(&batch)?.into_array(1)?; + let compare = spark_comparator(lists.as_ref(), value.as_ref())?; + let ordering = |i: usize| match (lists.is_null(i), value.is_null(0)) { + (false, false) => Some(Some(compare(i, 0))), + (true, true) => Some(None), + _ => None, + }; + let expr = spark_comparison( + Arc::clone(&a), + op, + Arc::clone(&lit), + &schema, + FloatOperands::Normalize, + )?; + assert!(expr.downcast_ref::().is_some() || !is_ordering(op)); + let expected: Vec<_> = (0..lists.len()) + .map(|i| spark_answer(op, ordering(i))) + .collect(); + assert_eq!(results(&expr, &batch), expected, "a {op} row {row}"); + let reversed = spark_comparison( + Arc::clone(&lit), + op.swap().unwrap(), + Arc::clone(&a), + &schema, + FloatOperands::Normalize, + )?; + assert_eq!( + results(&reversed, &batch), + expected, + "row {row} swapped {op} a" + ); + let both = spark_comparison( + Arc::clone(&lit), + op, + Arc::clone(&lit), + &schema, + FloatOperands::Normalize, + )?; + let ColumnarValue::Scalar(ScalarValue::Boolean(actual)) = both.evaluate(&batch)? + else { + panic!("two literals must give a Boolean scalar"); + }; + let ordering = if value.is_null(0) { + Some(None) + } else { + Some(Some(std::cmp::Ordering::Equal)) + }; + assert_eq!(actual, spark_answer(op, ordering), "row {row} {op} itself"); + } + let expr = spark_comparison( + Arc::clone(&a), + op, + Arc::clone(&null), + &schema, + FloatOperands::Normalize, + )?; + let expected: Vec<_> = (0..lists.len()) + .map(|i| spark_answer(op, if lists.is_null(i) { Some(None) } else { None })) + .collect(); + assert_eq!(results(&expr, &batch), expected, "a {op} NULL"); + } + Ok(()) + } + + fn is_ordering(op: Operator) -> bool { + !matches!(op, Operator::Eq | Operator::NotEq) + } + + /// Ordering and null-safe comparisons of lists and structs match `spark_comparator`, which + /// compares float leaves in Spark's order without normalizing them. + #[test] + fn nested_float_operands_follow_spark_ordering() -> Result<()> { + use crate::float_semantics::spark_comparator; + let leaves = [ + Some(-0.0), + Some(0.0), + Some(f64::NAN), + Some(f64::from_bits(0xfff8_0000_0000_0000)), + Some(f64::INFINITY), + Some(1.0), + None, + ]; + // Lists of every pair of leaves, including lists of different lengths. + let mut left_lists = vec![]; + let mut right_lists = vec![]; + for &l in &leaves { + for &r in &leaves { + left_lists.push(Some(vec![l])); + right_lists.push(Some(vec![r])); + left_lists.push(Some(vec![Some(1.0), l])); + right_lists.push(Some(vec![Some(1.0), r, Some(2.0)])); + } + } + left_lists.push(None); + right_lists.push(Some(vec![])); + left_lists.push(None); + right_lists.push(None); + let list = |lists: &[Option>>]| -> ArrayRef { + Arc::new(ListArray::from_iter_primitive::( + lists.to_vec(), + )) + }; + let (left, right) = (list(&left_lists), list(&right_lists)); + let as_struct = |array: &ArrayRef| -> ArrayRef { + Arc::new(StructArray::new( + vec![Arc::new(Field::new("v", array.data_type().clone(), true))].into(), + vec![Arc::clone(array)], + None, + )) + }; + // Also from an offset, which shifts the outer null buffers that the comparison skips. + let sliced = |array: &ArrayRef| array.slice(3, array.len() - 3); + for (l, r) in [ + (Arc::clone(&left), Arc::clone(&right)), + (as_struct(&left), as_struct(&right)), + (sliced(&left), sliced(&right)), + (sliced(&as_struct(&left)), sliced(&as_struct(&right))), + ] { + let compare = spark_comparator(l.as_ref(), r.as_ref())?; + let batch = batch(Arc::clone(&l), Arc::clone(&r)); + let (a, b) = columns(); + for op in COMPARISONS { + let expr = spark_comparison( + Arc::clone(&a), + op, + Arc::clone(&b), + batch.schema().as_ref(), + FloatOperands::Normalize, + )?; + let expected: Vec<_> = (0..l.len()) + .map(|i| { + let ordering = match (l.is_null(i), r.is_null(i)) { + (false, false) => Some(Some(compare(i, i))), + (true, true) => Some(None), + _ => None, + }; + spark_answer(op, ordering) + }) + .collect(); + assert_eq!(results(&expr, &batch), expected, "{} {op}", l.data_type()); + } + } + Ok(()) + } } diff --git a/native/spark-expr/src/bloom_filter/bloom_filter_might_contain.rs b/native/spark-expr/src/bloom_filter/bloom_filter_might_contain.rs index 66168444d9a..e60e6f462c1 100644 --- a/native/spark-expr/src/bloom_filter/bloom_filter_might_contain.rs +++ b/native/spark-expr/src/bloom_filter/bloom_filter_might_contain.rs @@ -57,6 +57,8 @@ fn evaluate_bloom_filter( ColumnarValue::Scalar(ScalarValue::Binary(v)) => { Ok(v.map(|v| SparkBloomFilter::from(v.as_slice()))) } + // An untyped NULL literal, as in `might_contain(null, 1L)` + ColumnarValue::Scalar(ScalarValue::Null) => Ok(None), _ => internal_err!("Bloom filter expression should be evaluated as a scalar binary value"), } } @@ -90,6 +92,10 @@ fn execute_might_contain( .and_then(|filter| optional_value.map(|value| filter.might_contain_long(value))); Ok(ColumnarValue::Scalar(ScalarValue::Boolean(result))) } + // An untyped NULL literal, as in `might_contain(bf, null)` + ColumnarValue::Scalar(ScalarValue::Null) => { + Ok(ColumnarValue::Scalar(ScalarValue::Boolean(None))) + } ColumnarValue::Array(values_array) => { let values = values_array .as_any() @@ -109,7 +115,7 @@ fn execute_might_contain( }) .unwrap_or_else(|| Ok(ColumnarValue::Scalar(ScalarValue::Boolean(None)))) } - _ => internal_err!("Expected Int64Array or Int64 Scalar as arguments"), + other => internal_err!("Expected Int64Array or Int64 Scalar as arguments, got {other:?}"), } } @@ -117,6 +123,7 @@ fn execute_might_contain( mod tests { use super::*; use arrow::array::BooleanArray; + use datafusion::physical_expr::expressions::Literal; fn assert_result_eq>>(result: ColumnarValue, expected: Vec) { let array = result.to_array(1).unwrap(); @@ -171,6 +178,24 @@ mod tests { assert_all_null(result); } + #[test] + fn test_execute_untyped_null_value() { + let mut filter = SparkBloomFilter::from((4, 64)); + filter.put_long(123); + + let args = [ColumnarValue::Scalar(ScalarValue::Null)]; + + let result = execute_might_contain(&Some(filter), &args).unwrap(); + assert_all_null(result); + } + + #[test] + fn test_untyped_null_bloom_filter() { + let bloom_filter_expr: Arc = Arc::new(Literal::new(ScalarValue::Null)); + let expr = BloomFilterMightContain::try_new(bloom_filter_expr).unwrap(); + assert!(expr.bloom_filter.is_none()); + } + #[test] fn test_execute_array() { let mut filter = SparkBloomFilter::from((4, 64)); diff --git a/native/spark-expr/src/comet_scalar_funcs.rs b/native/spark-expr/src/comet_scalar_funcs.rs index 0e4a3824e96..9cc1cadbd8d 100644 --- a/native/spark-expr/src/comet_scalar_funcs.rs +++ b/native/spark-expr/src/comet_scalar_funcs.rs @@ -34,9 +34,10 @@ use crate::{ spark_round, spark_rpad, spark_sequence, spark_to_time, spark_unhex, spark_unscaled_value, EvalMode, SparkArrayExtrema, SparkArrayPositionFunc, SparkArrayRemove, SparkArraySlice, SparkArraysOverlap, SparkContains, SparkDateDiff, SparkDateFromUnixDate, SparkDateTrunc, - SparkDayOfWeek, SparkFlatten, SparkIcebergBucket, SparkIcebergTemporalTransform, - SparkIcebergTruncate, SparkMakeDate, SparkMakeInterval, SparkMakeTime, SparkMapExtract, - SparkNextDay, SparkSecondsToTimestamp, SparkSizeFunc, SparkSortArray, SparkWeekDay, + SparkDayOfWeek, SparkFlatten, SparkFloatArrayContains, SparkIcebergBucket, + SparkIcebergTemporalTransform, SparkIcebergTruncate, SparkMakeDate, SparkMakeInterval, + SparkMakeTime, SparkMapExtract, SparkNextDay, SparkSecondsToTimestamp, SparkSizeFunc, + SparkSortArray, SparkWeekDay, }; use arrow::datatypes::DataType; use datafusion::common::{DataFusionError, Result as DataFusionResult}; @@ -342,6 +343,7 @@ fn all_scalar_functions() -> Vec> { Arc::new(ScalarUDF::new_from_impl(SparkArrayExtrema::new(false))), Arc::new(ScalarUDF::new_from_impl(SparkArrayPositionFunc::default())), Arc::new(ScalarUDF::new_from_impl(SparkArrayRemove::default())), + Arc::new(ScalarUDF::new_from_impl(SparkFloatArrayContains::default())), Arc::new(ScalarUDF::new_from_impl(SparkArraySlice::default())), Arc::new(ScalarUDF::new_from_impl(SparkArraysOverlap::default())), Arc::new(ScalarUDF::new_from_impl(SparkSortArray::default())), diff --git a/native/spark-expr/src/conditional_funcs/case_when.rs b/native/spark-expr/src/conditional_funcs/case_when.rs index b128b165093..74df26c977c 100644 --- a/native/spark-expr/src/conditional_funcs/case_when.rs +++ b/native/spark-expr/src/conditional_funcs/case_when.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use crate::{Cast, EvalMode, IfExpr, NormalizeNaNAndZero, SparkCastOptions}; +use crate::{Cast, EvalMode, IfExpr, NormalizeNaNAndZero, SparkCastOptions, SparkComparison}; use arrow::array::{ downcast_primitive, Array, ArrayRef, AsArray, BooleanArray, GenericByteArray, PrimitiveArray, }; @@ -152,7 +152,7 @@ fn coerce_branch( if data_type == common_type { return expr; } - let cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC"); Arc::new(Cast::new( expr, common_type.clone(), @@ -396,6 +396,8 @@ fn is_infallible(expr: &Arc, input_schema: &Schema) -> bool { } let node_is_infallible = if let Some(binary) = expr.downcast_ref::() { binary_is_infallible(binary, input_schema) + } else if let Some(comparison) = expr.downcast_ref::() { + comparison.is_infallible(input_schema) } else if let Some(cast) = expr.downcast_ref::() { cast.is_infallible(input_schema) } else if let Some(normalize) = expr.downcast_ref::() { @@ -836,6 +838,7 @@ fn merge_nulls(num_rows: usize, branches: &[Branch]) -> Option { #[cfg(test)] mod tests { use super::*; + use crate::{spark_comparison, FloatOperands}; use arrow::array::{ BinaryArray, Int32Array, Int64Array, LargeStringArray, StringArray, StructArray, TimestampMicrosecondArray, @@ -1137,24 +1140,40 @@ mod tests { #[test] fn infallible_expressions() { + let list_of = |item: DataType| DataType::List(Arc::new(Field::new("item", item, true))); let schema = Schema::new(vec![ Field::new("i", DataType::Int64, true), Field::new("j", DataType::Int32, true), Field::new("d", DataType::Float64, true), + Field::new("e", DataType::Float64, true), Field::new("s", DataType::Utf8, true), Field::new("dec", DataType::Decimal128(10, 2), true), + Field::new("l", list_of(DataType::Float64), true), + Field::new("m", list_of(DataType::Float64), true), + Field::new( + "dictionary", + list_of(DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Float64), + )), + true, + ), ]); let c = |name: &str| col(name, &schema).unwrap(); let cast = |e: Arc, to: DataType| -> Arc { Arc::new(Cast::new( e, to, - SparkCastOptions::new_without_timezone(EvalMode::Ansi, false), + SparkCastOptions::new_without_timezone(EvalMode::Ansi), None, None, )) }; let infallible = |e: Arc| is_infallible(&e, &schema); + // A comparison as the planner builds it, which follows Spark's ordering for floats + let compare = |left: &str, op: Operator, right: Arc| { + spark_comparison(c(left), op, right, &schema, FloatOperands::Normalize).unwrap() + }; assert!(infallible(c("i"))); assert!(infallible(lit("x"))); @@ -1181,6 +1200,13 @@ mod tests { DataType::Float64, c("d") )))); + for op in [Operator::Lt, Operator::Eq, Operator::IsNotDistinctFrom] { + assert!(infallible(compare("d", op, lit(1.5))), "d {op} 1.5"); + assert!(infallible(compare("d", op, c("e"))), "d {op} e"); + } + for op in [Operator::Lt, Operator::IsNotDistinctFrom] { + assert!(infallible(compare("l", op, c("m"))), "l {op} m"); + } // Only a boolean can be negated, and only a float normalized assert!(!infallible(Arc::new(NotExpr::new(c("i"))))); @@ -1188,6 +1214,8 @@ mod tests { DataType::Int64, c("i") )))); + // A nested comparison fails on a dictionary-encoded leaf against a plain one + assert!(!infallible(compare("l", Operator::Lt, c("dictionary")))); // Integer division and remainder fail on a zero divisor assert!(!infallible(binary(c("i"), Operator::Divide, c("i")))); assert!(!infallible(binary(c("i"), Operator::Modulo, c("i")))); diff --git a/native/spark-expr/src/conversion_funcs/boolean.rs b/native/spark-expr/src/conversion_funcs/boolean.rs index 1db2746ce24..7e9dffd3307 100644 --- a/native/spark-expr/src/conversion_funcs/boolean.rs +++ b/native/spark-expr/src/conversion_funcs/boolean.rs @@ -67,7 +67,7 @@ mod tests { } fn test_input_spark_opts() -> SparkCastOptions { - SparkCastOptions::new(EvalMode::Legacy, "Asia/Kolkata", false) + SparkCastOptions::new(EvalMode::Legacy, "Asia/Kolkata") } #[test] diff --git a/native/spark-expr/src/conversion_funcs/cast.rs b/native/spark-expr/src/conversion_funcs/cast.rs index 26659f7489f..30b3fe997a7 100644 --- a/native/spark-expr/src/conversion_funcs/cast.rs +++ b/native/spark-expr/src/conversion_funcs/cast.rs @@ -149,8 +149,6 @@ pub struct SparkCastOptions { /// session local timezone by an analyzer in Spark. // TODO we should change timezone to Tz to avoid repeated parsing pub timezone: String, - /// Allow casts that are supported but not guaranteed to be 100% compatible - pub allow_incompat: bool, /// True when running against Spark 4.0+. Enables version-specific cast behaviour /// such as the handling of leading whitespace before T-prefixed time-only strings. pub is_spark4_plus: bool, @@ -166,11 +164,10 @@ pub struct SparkCastOptions { } impl SparkCastOptions { - pub fn new(eval_mode: EvalMode, timezone: &str, allow_incompat: bool) -> Self { + pub fn new(eval_mode: EvalMode, timezone: &str) -> Self { Self { eval_mode, timezone: timezone.to_string(), - allow_incompat, is_spark4_plus: false, allow_cast_unsigned_ints: false, is_adapting_schema: false, @@ -179,11 +176,10 @@ impl SparkCastOptions { } } - pub fn new_without_timezone(eval_mode: EvalMode, allow_incompat: bool) -> Self { + pub fn new_without_timezone(eval_mode: EvalMode) -> Self { Self { eval_mode, timezone: "".to_string(), - allow_incompat, is_spark4_plus: false, allow_cast_unsigned_ints: false, is_adapting_schema: false, @@ -192,15 +188,10 @@ impl SparkCastOptions { } } - pub fn new_with_version( - eval_mode: EvalMode, - timezone: &str, - allow_incompat: bool, - is_spark4_plus: bool, - ) -> Self { + pub fn new_with_version(eval_mode: EvalMode, timezone: &str, is_spark4_plus: bool) -> Self { Self { is_spark4_plus, - ..Self::new(eval_mode, timezone, allow_incompat) + ..Self::new(eval_mode, timezone) } } } @@ -941,7 +932,7 @@ mod tests { let error = cast_array( Arc::new(StringArray::from(vec!["a"])), &DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), - &SparkCastOptions::new(EvalMode::Legacy, "UTC", false), + &SparkCastOptions::new(EvalMode::Legacy, "UTC"), ) .unwrap_err(); @@ -979,7 +970,7 @@ mod tests { let error = cast_array( input, &output_type, - &SparkCastOptions::new_without_timezone(EvalMode::Ansi, false), + &SparkCastOptions::new_without_timezone(EvalMode::Ansi), ) .unwrap_err(); @@ -1014,7 +1005,7 @@ mod tests { Some(&[0xEDu8, 0xA0, 0x80][..]), ]); // binary_output_style defaults to None, i.e. the plain (non-ToPrettyString) cast path. - let cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC"); let result = cast_binary_to_string::(&input, &cast_options).unwrap(); @@ -1038,7 +1029,7 @@ mod tests { Some("abc".as_bytes()), Some(&[0xEDu8, 0xA0, 0x80][..]), ]); - let mut cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let mut cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC"); cast_options.binary_output_style = Some(BinaryOutputStyle::Utf8); let result = cast_binary_to_string::(&input, &cast_options).unwrap(); @@ -1060,7 +1051,7 @@ mod tests { Some(b"hi".as_slice()), ])); let cast = |style: Option| { - let mut options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let mut options = SparkCastOptions::new(EvalMode::Legacy, "UTC"); options.binary_output_style = style; let result = spark_cast( ColumnarValue::Array(Arc::clone(&input)), @@ -1124,7 +1115,7 @@ mod tests { None, Some("héllo".as_bytes()), ])); - let options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let options = SparkCastOptions::new(EvalMode::Legacy, "UTC"); let result = spark_cast( ColumnarValue::Array(Arc::clone(&input)), &DataType::Utf8, @@ -1143,7 +1134,7 @@ mod tests { fn test_cast_unsupported_timestamp_to_date() { // Since datafusion uses chrono::Datetime internally not all dates representable by TimestampMicrosecondType are supported let timestamps: PrimitiveArray = vec![i64::MAX].into(); - let cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC", false); + let cast_options = SparkCastOptions::new(EvalMode::Legacy, "UTC"); let result = cast_array( Arc::new(timestamps.with_timezone("Europe/Copenhagen")), &DataType::Date32, @@ -1155,7 +1146,7 @@ mod tests { #[test] fn test_cast_invalid_timezone() { let timestamps: PrimitiveArray = vec![i64::MAX].into(); - let cast_options = SparkCastOptions::new(EvalMode::Legacy, "Not a valid timezone", false); + let cast_options = SparkCastOptions::new(EvalMode::Legacy, "Not a valid timezone"); let result = cast_array( Arc::new(timestamps.with_timezone("Europe/Copenhagen")), &DataType::Date32, @@ -1181,7 +1172,7 @@ mod tests { let string_array = cast_array( c, &DataType::Utf8, - &SparkCastOptions::new(EvalMode::Legacy, "UTC", false), + &SparkCastOptions::new(EvalMode::Legacy, "UTC"), ) .unwrap(); let string_array = string_array.as_string::(); @@ -1215,7 +1206,7 @@ mod tests { let cast_array = spark_cast( ColumnarValue::Array(c), &DataType::Struct(fields), - &SparkCastOptions::new(EvalMode::Legacy, "UTC", false), + &SparkCastOptions::new(EvalMode::Legacy, "UTC"), ) .unwrap(); if let ColumnarValue::Array(cast_array) = cast_array { @@ -1253,7 +1244,7 @@ mod tests { let result = spark_cast( ColumnarValue::Array(outer), &DataType::Struct(to_fields), - &SparkCastOptions::new(EvalMode::Ansi, "UTC", false), + &SparkCastOptions::new(EvalMode::Ansi, "UTC"), ); assert!(result.is_err()); @@ -1278,7 +1269,7 @@ mod tests { let cast_array = spark_cast( ColumnarValue::Array(c), &DataType::Struct(fields), - &SparkCastOptions::new(EvalMode::Legacy, "UTC", false), + &SparkCastOptions::new(EvalMode::Legacy, "UTC"), ) .unwrap(); if let ColumnarValue::Array(cast_array) = cast_array { @@ -1304,11 +1295,9 @@ mod tests { Arc::new(values_array), None, )); - let string_array = cast_array_to_string( - &list_array, - &SparkCastOptions::new(EvalMode::Legacy, "UTC", false), - ) - .unwrap(); + let string_array = + cast_array_to_string(&list_array, &SparkCastOptions::new(EvalMode::Legacy, "UTC")) + .unwrap(); let string_array = string_array.as_string::(); assert_eq!(r#"[a, b, c]"#, string_array.value(0)); assert_eq!(r#"[a, null]"#, string_array.value(1)); @@ -1327,11 +1316,9 @@ mod tests { Arc::new(values_array), None, )); - let string_array = cast_array_to_string( - &list_array, - &SparkCastOptions::new(EvalMode::Legacy, "UTC", false), - ) - .unwrap(); + let string_array = + cast_array_to_string(&list_array, &SparkCastOptions::new(EvalMode::Legacy, "UTC")) + .unwrap(); let string_array = string_array.as_string::(); assert_eq!(r#"[1, 2, 3]"#, string_array.value(0)); assert_eq!(r#"[1, null]"#, string_array.value(1)); @@ -1354,7 +1341,7 @@ mod tests { let to_array = cast_array( from_array, &to_type, - &SparkCastOptions::new(EvalMode::Legacy, "UTC", false), + &SparkCastOptions::new(EvalMode::Legacy, "UTC"), ) .unwrap(); @@ -1368,7 +1355,7 @@ mod tests { assert!(values.iter().all(|value| value.is_none())); } fn legacy_opts() -> SparkCastOptions { - SparkCastOptions::new(EvalMode::Legacy, "UTC", false) + SparkCastOptions::new(EvalMode::Legacy, "UTC") } /// Build a `Map` MapArray (Parquet-style "key_value" field names). @@ -1532,7 +1519,7 @@ mod tests { let casted = cast_array( map_array, &to_type, - &SparkCastOptions::new(EvalMode::Legacy, "UTC", false), + &SparkCastOptions::new(EvalMode::Legacy, "UTC"), ) .unwrap(); diff --git a/native/spark-expr/src/conversion_funcs/numeric.rs b/native/spark-expr/src/conversion_funcs/numeric.rs index f7fbae54f98..aad9006937c 100644 --- a/native/spark-expr/src/conversion_funcs/numeric.rs +++ b/native/spark-expr/src/conversion_funcs/numeric.rs @@ -2291,7 +2291,7 @@ mod tests { ); for eval_mode in [EvalMode::Legacy, EvalMode::Ansi, EvalMode::Try] { - let options = SparkCastOptions::new(eval_mode, "UTC", false); + let options = SparkCastOptions::new(eval_mode, "UTC"); let doubles = cast_array(Arc::clone(&decimals_38_18), &DataType::Float64, &options).unwrap(); diff --git a/native/spark-expr/src/conversion_funcs/string.rs b/native/spark-expr/src/conversion_funcs/string.rs index adedd3ba0d2..f278bafda68 100644 --- a/native/spark-expr/src/conversion_funcs/string.rs +++ b/native/spark-expr/src/conversion_funcs/string.rs @@ -2491,7 +2491,7 @@ mod tests { for (position, input, expect_value) in cases { for eval_mode in [EvalMode::Legacy, EvalMode::Try, EvalMode::Ansi] { let array: ArrayRef = Arc::new(StringArray::from(vec![Some(input.as_str())])); - let options = SparkCastOptions::new(eval_mode, "UTC", false); + let options = SparkCastOptions::new(eval_mode, "UTC"); let result = cast_array(array, to_type, &options); let context = format!("cast {input:?} ({position}) to {to_type} in {eval_mode:?}"); if expect_value { @@ -2736,7 +2736,7 @@ mod tests { let timezone = "UTC".to_string(); // test casting string dictionary array to timestamp array - let cast_options = SparkCastOptions::new(EvalMode::Legacy, &timezone, false); + let cast_options = SparkCastOptions::new(EvalMode::Legacy, &timezone); let result = cast_array( dict_array, &DataType::Timestamp(TimeUnit::Microsecond, Some(timezone.clone().into())), diff --git a/native/spark-expr/src/conversion_funcs/temporal.rs b/native/spark-expr/src/conversion_funcs/temporal.rs index 31e63644fd5..15926047f73 100644 --- a/native/spark-expr/src/conversion_funcs/temporal.rs +++ b/native/spark-expr/src/conversion_funcs/temporal.rs @@ -121,7 +121,7 @@ mod tests { let target_tz: Option> = Some("UTC".into()); let result = cast_date_to_timestamp( &dates, - &SparkCastOptions::new(EvalMode::Legacy, "UTC", false), + &SparkCastOptions::new(EvalMode::Legacy, "UTC"), &target_tz, ) .unwrap(); @@ -134,7 +134,7 @@ mod tests { // validate LA timezone (follows Daylight savings) let result = cast_date_to_timestamp( &dates, - &SparkCastOptions::new(EvalMode::Legacy, "America/Los_Angeles", false), + &SparkCastOptions::new(EvalMode::Legacy, "America/Los_Angeles"), &target_tz, ) .unwrap(); @@ -148,7 +148,7 @@ mod tests { // Phoenix timezone (does not follow Daylight savings) let result = cast_date_to_timestamp( &dates, - &SparkCastOptions::new(EvalMode::Legacy, "America/Phoenix", false), + &SparkCastOptions::new(EvalMode::Legacy, "America/Phoenix"), &target_tz, ) .unwrap(); @@ -189,7 +189,7 @@ mod tests { ] { let result = cast_date_to_timestamp( &dates, - &SparkCastOptions::new(EvalMode::Legacy, tz, false), + &SparkCastOptions::new(EvalMode::Legacy, tz), &ntz_target, ) .unwrap(); @@ -215,7 +215,7 @@ mod tests { assert!( cast_date_to_timestamp( &dates, - &SparkCastOptions::new(EvalMode::Legacy, "UTC", false), + &SparkCastOptions::new(EvalMode::Legacy, "UTC"), &ntz_target, ) .is_err(), @@ -240,7 +240,7 @@ mod tests { Some(-106_751_992), ]; for mode in [EvalMode::Legacy, EvalMode::Ansi, EvalMode::Try] { - let options = SparkCastOptions::new(mode, "America/Los_Angeles", false); + let options = SparkCastOptions::new(mode, "America/Los_Angeles"); let result = spark_cast( ColumnarValue::Array(Arc::new(Date32Array::from(days.clone()))), &target, diff --git a/native/spark-expr/src/csv_funcs/to_csv.rs b/native/spark-expr/src/csv_funcs/to_csv.rs index 01fdc901cb7..d79b5cbca3d 100644 --- a/native/spark-expr/src/csv_funcs/to_csv.rs +++ b/native/spark-expr/src/csv_funcs/to_csv.rs @@ -86,7 +86,7 @@ impl PhysicalExpr for ToCsv { fn evaluate(&self, batch: &RecordBatch) -> Result { let input_array = self.expr.evaluate(batch)?.into_array(batch.num_rows())?; - let mut cast_options = SparkCastOptions::new(EvalMode::Legacy, &self.timezone, false); + let mut cast_options = SparkCastOptions::new(EvalMode::Legacy, &self.timezone); cast_options.null_string = self.csv_write_options.null_value.clone(); let struct_array = as_struct_array(&input_array); diff --git a/native/spark-expr/src/datetime_funcs/timestamp_trunc.rs b/native/spark-expr/src/datetime_funcs/timestamp_trunc.rs index 91a83b77e50..e5c07ebbb3b 100644 --- a/native/spark-expr/src/datetime_funcs/timestamp_trunc.rs +++ b/native/spark-expr/src/datetime_funcs/timestamp_trunc.rs @@ -16,7 +16,7 @@ // under the License. use crate::utils::array_with_timezone; -use arrow::array::ArrayRef; +use arrow::array::{new_null_array, ArrayRef}; use arrow::compute::cast; use arrow::datatypes::{DataType, Schema, TimeUnit::Microsecond}; use arrow::record_batch::RecordBatch; @@ -41,6 +41,9 @@ pub struct TimestampTruncExpr { /// `Arc` so it can be cheaply cloned onto Arrow `Timestamp` data types without /// reallocating, and parsed once into a `chrono::TimeZone` per batch. timezone: Arc, + /// Whether SECOND and MILLISECOND truncation wraps below the smallest timestamp, as Spark + /// did before 4.2.0, instead of raising `long overflow` as 4.2.0 and later do (SPARK-56663). + wrap_second_millisecond_overflow: bool, } impl Hash for TimestampTruncExpr { @@ -48,6 +51,7 @@ impl Hash for TimestampTruncExpr { self.child.hash(state); self.format.hash(state); self.timezone.hash(state); + self.wrap_second_millisecond_overflow.hash(state); } } impl PartialEq for TimestampTruncExpr { @@ -55,6 +59,7 @@ impl PartialEq for TimestampTruncExpr { self.child.eq(&other.child) && self.format.eq(&other.format) && self.timezone.eq(&other.timezone) + && self.wrap_second_millisecond_overflow == other.wrap_second_millisecond_overflow } } @@ -63,11 +68,13 @@ impl TimestampTruncExpr { child: Arc, format: Arc, timezone: String, + wrap_second_millisecond_overflow: bool, ) -> Self { TimestampTruncExpr { child, format, timezone: Arc::from(timezone), + wrap_second_millisecond_overflow, } } } @@ -100,6 +107,7 @@ impl PhysicalExpr for TimestampTruncExpr { let format = self.format.evaluate(batch)?; let output_type = output_type(×tamp.data_type()); let tz = &self.timezone; + let wrap = self.wrap_second_millisecond_overflow; let resolve_tz = |ts: ArrayRef| -> datafusion::common::Result { // For TimestampNTZ (Timestamp(Microsecond, None)), skip timezone conversion. // NTZ values are timezone-independent and truncation should operate directly on the @@ -121,18 +129,31 @@ impl PhysicalExpr for TimestampTruncExpr { }; match (timestamp, format) { (ColumnarValue::Array(ts), ColumnarValue::Scalar(Utf8(Some(format)))) => { - let result = timestamp_trunc_dyn(&resolve_tz(ts)?, format)?; + let result = timestamp_trunc_dyn(&resolve_tz(ts)?, format, wrap)?; Ok(ColumnarValue::Array(relabel(result)?)) } (ColumnarValue::Array(ts), ColumnarValue::Array(formats)) => { - let result = timestamp_trunc_array_fmt_dyn(&resolve_tz(ts)?, &formats)?; + let result = timestamp_trunc_array_fmt_dyn(&resolve_tz(ts)?, &formats, wrap)?; Ok(ColumnarValue::Array(relabel(result)?)) } (ColumnarValue::Scalar(ts_scalar), ColumnarValue::Scalar(Utf8(Some(format)))) => { - let result = timestamp_trunc_dyn(&resolve_tz(ts_scalar.to_array()?)?, format)?; + let result = + timestamp_trunc_dyn(&resolve_tz(ts_scalar.to_array()?)?, format, wrap)?; let scalar = ScalarValue::try_from_array(&relabel(result)?, 0)?; Ok(ColumnarValue::Scalar(scalar)) } + (ColumnarValue::Scalar(ts_scalar), ColumnarValue::Array(formats)) => { + let ts = ts_scalar.to_array_of_size(formats.len())?; + let result = timestamp_trunc_array_fmt_dyn(&resolve_tz(ts)?, &formats, wrap)?; + Ok(ColumnarValue::Array(relabel(result)?)) + } + // A NULL format gives NULL, as in Spark. + (ColumnarValue::Array(ts), ColumnarValue::Scalar(Utf8(None))) => { + Ok(ColumnarValue::Array(new_null_array(&output_type, ts.len()))) + } + (ColumnarValue::Scalar(_), ColumnarValue::Scalar(Utf8(None))) => { + Ok(ColumnarValue::Scalar(ScalarValue::try_from(&output_type)?)) + } _ => Err(DataFusionError::Execution( "Invalid input to function TimestampTrunc. \ Expected (Timestamp, Utf8)" @@ -153,6 +174,7 @@ impl PhysicalExpr for TimestampTruncExpr { Arc::clone(&children[0]), Arc::clone(&self.format), self.timezone.to_string(), + self.wrap_second_millisecond_overflow, ))) } } @@ -174,7 +196,9 @@ fn output_type(input: &DataType) -> DataType { #[cfg(test)] mod tests { use super::*; - use arrow::array::{Array, AsArray, DictionaryArray, Int32Array, TimestampMicrosecondArray}; + use arrow::array::{ + Array, AsArray, DictionaryArray, Int32Array, StringArray, TimestampMicrosecondArray, + }; use arrow::datatypes::{Field, Int32Type, TimestampMicrosecondType}; use datafusion::physical_expr::expressions::{Column, Literal}; @@ -195,6 +219,7 @@ mod tests { Arc::new(Column::new("ts", 0)), Arc::new(Literal::new(Utf8(Some("HOUR".to_string())))), "Asia/Kolkata".to_string(), + false, ); let declared = expr.data_type(&schema).unwrap(); let ColumnarValue::Array(result) = expr.evaluate(&batch).unwrap() else { @@ -230,4 +255,68 @@ mod tests { HOUR_IN_KOLKATA ); } + + /// Evaluates `date_trunc` in America/Los_Angeles with the given timestamp and format inputs. + fn trunc_in_los_angeles( + timestamp: Arc, + format: Arc, + batch: &RecordBatch, + ) -> ColumnarValue { + TimestampTruncExpr::new(timestamp, format, "America/Los_Angeles".to_string(), false) + .evaluate(batch) + .unwrap() + } + + #[test] + fn literal_timestamp_with_a_format_column() { + // 2024-11-03 01:30 PDT, inside the overlap. DAY keeps the input's offset, YEAR does not + // depend on it, and a NULL format gives NULL. + let micros = 1_730_622_600_000_000; + let formats = StringArray::from(vec![Some("DAY"), Some("YEAR"), None]); + let schema = Schema::new(vec![Field::new("fmt", DataType::Utf8, true)]); + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(formats)]).unwrap(); + let timestamp = Literal::new(ScalarValue::TimestampMicrosecond( + Some(micros), + Some("UTC".into()), + )); + let ColumnarValue::Array(result) = + trunc_in_los_angeles(Arc::new(timestamp), Arc::new(Column::new("fmt", 0)), &batch) + else { + panic!("expected an array"); + }; + assert_eq!(result.data_type(), &utc_timestamp()); + let result = result.as_primitive::(); + // 2024-11-03 00:00 PDT and 2024-01-01 00:00 PST. + assert_eq!(result.value(0), 1_730_617_200_000_000); + assert_eq!(result.value(1), 1_704_096_000_000_000); + assert!(result.is_null(2)); + } + + #[test] + fn null_literal_format_gives_null() { + let input = TimestampMicrosecondArray::from(vec![Some(MICROS), None]).with_timezone("UTC"); + let schema = Schema::new(vec![Field::new("ts", utc_timestamp(), true)]); + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(input)]).unwrap(); + let null_format = || Arc::new(Literal::new(Utf8(None))); + let ColumnarValue::Array(result) = + trunc_in_los_angeles(Arc::new(Column::new("ts", 0)), null_format(), &batch) + else { + panic!("expected an array"); + }; + assert_eq!(result.data_type(), &utc_timestamp()); + assert_eq!(result.null_count(), 2); + + let timestamp = Arc::new(Literal::new(ScalarValue::TimestampMicrosecond( + Some(MICROS), + Some("UTC".into()), + ))); + let ColumnarValue::Scalar(result) = trunc_in_los_angeles(timestamp, null_format(), &batch) + else { + panic!("expected a scalar"); + }; + assert_eq!( + result, + ScalarValue::TimestampMicrosecond(None, Some("UTC".into())) + ); + } } diff --git a/native/spark-expr/src/float_semantics/compare.rs b/native/spark-expr/src/float_semantics/compare.rs index 24c01e1e9df..a9d314cb730 100644 --- a/native/spark-expr/src/float_semantics/compare.rs +++ b/native/spark-expr/src/float_semantics/compare.rs @@ -17,6 +17,7 @@ use super::{compare_floats, has_float_leaf}; use arrow::array::{make_comparator, Array, AsArray, DynComparator, OffsetSizeTrait}; +use arrow::buffer::NullBuffer; use arrow::compute::SortOptions; use arrow::datatypes::{ArrowPrimitiveType, DataType, Float32Type, Float64Type}; use datafusion::common::{internal_err, DFSchema, Result}; @@ -32,7 +33,7 @@ use std::ops::Range; /// arrays must have the same type, ignoring field names and nullability. pub fn spark_comparator(left: &dyn Array, right: &dyn Array) -> Result { check_types(left, right)?; - comparator(left, right, false) + comparator(left, right, false, true) } /// Builds a test of whether `left[i]` equals `right[j]` in the ordering of [`spark_comparator`]. @@ -42,7 +43,30 @@ pub fn spark_equality( right: &dyn Array, ) -> Result bool + Send + Sync>> { check_types(left, right)?; - let compare = comparator(left, right, true)?; + let compare = comparator(left, right, true, true)?; + Ok(Box::new(move |i, j| compare(i, j).is_eq())) +} + +/// [`spark_comparator`] for a caller that applies the top-level nulls of `left` and `right` +/// itself, as SQL comparisons do. A null list or struct is compared by the values under it instead +/// of sorting first, which saves testing both null buffers on every call. Nulls inside a list or +/// struct still sort first. +pub(crate) fn spark_comparator_ignoring_nulls( + left: &dyn Array, + right: &dyn Array, +) -> Result { + check_types(left, right)?; + comparator(left, right, false, false) +} + +/// [`spark_equality`] for a caller that applies the top-level nulls itself. See +/// [`spark_comparator_ignoring_nulls`]. +pub(crate) fn spark_equality_ignoring_nulls( + left: &dyn Array, + right: &dyn Array, +) -> Result bool + Send + Sync>> { + check_types(left, right)?; + let compare = comparator(left, right, true, false)?; Ok(Box::new(move |i, j| compare(i, j).is_eq())) } @@ -58,8 +82,15 @@ fn check_types(left: &dyn Array, right: &dyn Array) -> Result<()> { } } -/// With `equality` set, the comparator only has to tell equal from unequal values. -fn comparator(left: &dyn Array, right: &dyn Array, equality: bool) -> Result { +/// With `equality` set, the comparator only has to tell equal from unequal values. With `nulls` +/// unset, it ignores the nulls of `left` and `right` themselves, though not those of their +/// children. +fn comparator( + left: &dyn Array, + right: &dyn Array, + equality: bool, + nulls: bool, +) -> Result { if !has_float_leaf(left.data_type()) { let options = SortOptions { descending: false, @@ -68,38 +99,51 @@ fn comparator(left: &dyn Array, right: &dyn Array, equality: bool) -> Result Ok(float_comparator::(left, right)), - (DataType::Float64, DataType::Float64) => Ok(float_comparator::(left, right)), - (DataType::List(_), DataType::List(_)) => list_comparator::(left, right, equality), + (DataType::Float32, DataType::Float32) => { + Ok(float_comparator::(left, right, nulls)) + } + (DataType::Float64, DataType::Float64) => { + Ok(float_comparator::(left, right, nulls)) + } + (DataType::List(_), DataType::List(_)) => { + list_comparator::(left, right, equality, nulls) + } (DataType::LargeList(_), DataType::LargeList(_)) => { - list_comparator::(left, right, equality) + list_comparator::(left, right, equality, nulls) } (DataType::FixedSizeList(_, _), DataType::FixedSizeList(_, _)) => { - fixed_size_list_comparator(left, right, equality) + fixed_size_list_comparator(left, right, equality, nulls) + } + (DataType::Struct(_), DataType::Struct(_)) => { + struct_comparator(left, right, equality, nulls) } - (DataType::Struct(_), DataType::Struct(_)) => struct_comparator(left, right, equality), (l, r) => internal_err!("Unsupported types for Spark comparison: {l} and {r}"), } } -fn float_comparator(left: &dyn Array, right: &dyn Array) -> DynComparator +fn float_comparator( + left: &dyn Array, + right: &dyn Array, + nulls: bool, +) -> DynComparator where T::Native: Float, { let l = left.as_primitive::().values().clone(); let r = right.as_primitive::().values().clone(); - nulls_first(left, right, move |i, j| compare_floats(l[i], r[j])) + nulls_first(left, right, nulls, move |i, j| compare_floats(l[i], r[j])) } fn list_comparator( left: &dyn Array, right: &dyn Array, equality: bool, + nulls: bool, ) -> Result { let (l, r) = (left.as_list::(), right.as_list::()); - let compare = comparator(l.values().as_ref(), r.values().as_ref(), equality)?; + let compare = comparator(l.values().as_ref(), r.values().as_ref(), equality, true)?; let (l, r) = (l.offsets().clone(), r.offsets().clone()); - Ok(nulls_first(left, right, move |i, j| { + Ok(nulls_first(left, right, nulls, move |i, j| { let left = l[i].as_usize()..l[i + 1].as_usize(); let right = r[j].as_usize()..r[j + 1].as_usize(); lexicographic(&compare, left, right, equality) @@ -110,11 +154,12 @@ fn fixed_size_list_comparator( left: &dyn Array, right: &dyn Array, equality: bool, + nulls: bool, ) -> Result { let (l, r) = (left.as_fixed_size_list(), right.as_fixed_size_list()); - let compare = comparator(l.values().as_ref(), r.values().as_ref(), equality)?; + let compare = comparator(l.values().as_ref(), r.values().as_ref(), equality, true)?; let (l, r) = (l.value_length() as usize, r.value_length() as usize); - Ok(nulls_first(left, right, move |i, j| { + Ok(nulls_first(left, right, nulls, move |i, j| { lexicographic(&compare, i * l..(i + 1) * l, j * r..(j + 1) * r, equality) })) } @@ -137,15 +182,20 @@ fn lexicographic( .unwrap_or(lengths) } -fn struct_comparator(left: &dyn Array, right: &dyn Array, equality: bool) -> Result { +fn struct_comparator( + left: &dyn Array, + right: &dyn Array, + equality: bool, + nulls: bool, +) -> Result { let fields = left .as_struct() .columns() .iter() .zip(right.as_struct().columns()) - .map(|(l, r)| comparator(l.as_ref(), r.as_ref(), equality)) + .map(|(l, r)| comparator(l.as_ref(), r.as_ref(), equality, true)) .collect::>>()?; - Ok(nulls_first(left, right, move |i, j| { + Ok(nulls_first(left, right, nulls, move |i, j| { fields .iter() .map(|compare| compare(i, j)) @@ -154,18 +204,19 @@ fn struct_comparator(left: &dyn Array, right: &dyn Array, equality: bool) -> Res })) } -/// Orders nulls before values. A null slot never reaches `compare`, because the child values -/// under a null list or struct can be anything. +/// Orders nulls before values, unless `apply` is unset. A null slot then never reaches `compare`, +/// because the child values under a null list or struct can be anything. Without `apply`, the +/// caller masks the result for null slots: `compare` reads the values under them, which are in +/// bounds in any valid array. fn nulls_first( left: &dyn Array, right: &dyn Array, + apply: bool, compare: impl Fn(usize, usize) -> Ordering + Send + Sync + 'static, ) -> DynComparator { - let left = left.nulls().filter(|nulls| nulls.null_count() > 0).cloned(); - let right = right - .nulls() - .filter(|nulls| nulls.null_count() > 0) - .cloned(); + let present = |nulls: &&NullBuffer| apply && nulls.null_count() > 0; + let left = left.nulls().filter(present).cloned(); + let right = right.nulls().filter(present).cloned(); if left.is_none() && right.is_none() { return Box::new(compare); } @@ -311,6 +362,50 @@ mod tests { Ok(()) } + /// Ignoring the top-level nulls compares a null list or struct by the values under it, and + /// leaves every other slot, and nulls further down, as `spark_comparator` orders them. + #[test] + fn ignoring_nulls_compares_the_values_under_a_null_slot() -> Result<()> { + let field = Arc::new(Field::new("item", DataType::Float64, true)); + // The second row of each side is null, over `[-0.0]` and `[1.0]`, and the third holds an + // inner null. + let nulls = || Some(NullBuffer::from(vec![true, false, true])); + let list = |values: Vec>, nulls: Option| -> ArrayRef { + Arc::new(ListArray::new( + Arc::clone(&field), + OffsetBuffer::from_lengths([1, 1, 1]), + Arc::new(Float64Array::from(values)), + nulls, + )) + }; + let left = vec![Some(0.0), Some(-0.0), None]; + let right = vec![Some(-0.0), Some(1.0), Some(f64::NAN)]; + let struct_of = |values: Vec>| -> ArrayRef { + let list = list(values, None); + Arc::new(StructArray::new( + vec![Arc::new(Field::new("v", list.data_type().clone(), true))].into(), + vec![list], + nulls(), + )) + }; + for (l, r) in [ + (list(left.clone(), nulls()), list(right.clone(), nulls())), + (struct_of(left), struct_of(right)), + ] { + let masked = spark_comparator(l.as_ref(), r.as_ref())?; + let unmasked = spark_comparator_ignoring_nulls(l.as_ref(), r.as_ref())?; + let equal = spark_equality_ignoring_nulls(l.as_ref(), r.as_ref())?; + assert_eq!(masked(1, 1), Ordering::Equal); + assert_eq!(unmasked(1, 1), Ordering::Less, "{}", l.data_type()); + assert!(!equal(1, 1)); + for row in [0, 2] { + assert_eq!(unmasked(row, row), masked(row, row), "{}", l.data_type()); + assert_eq!(equal(row, row), masked(row, row).is_eq()); + } + } + Ok(()) + } + #[test] fn struct_fields_without_floats_break_ties() -> Result<()> { let structs = |floats: Vec, ints: Vec| -> ArrayRef { diff --git a/native/spark-expr/src/float_semantics/kernels.rs b/native/spark-expr/src/float_semantics/kernels.rs new file mode 100644 index 00000000000..aa04f7034bb --- /dev/null +++ b/native/spark-expr/src/float_semantics/kernels.rs @@ -0,0 +1,413 @@ +// 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. + +//! Comparison kernels for Float32 and Float64 arrays in Spark's SQL ordering, which read the +//! operands' value buffers as they are instead of normalizing copies of them first. +//! +//! Each operator is one test over a pair of values, written with `&` and `|` rather than `&&` and +//! `||` so that it stays free of branches, and the loops pack 64 results into a word at a time, +//! which the compiler vectorizes. + +use arrow::array::{Array, AsArray, BooleanArray, PrimitiveArray}; +use arrow::buffer::{BooleanBuffer, MutableBuffer, NullBuffer}; +use arrow::datatypes::{ArrowPrimitiveType, DataType, Float32Type, Float64Type}; +use arrow::util::bit_util; +use datafusion::common::{internal_err, Result, ScalarValue}; +use datafusion::logical_expr::Operator; +use num::Float; + +/// `left op right` for two Float32 or two Float64 arrays of the same length, where `op` is one of +/// `=`, `<>`, `<`, `<=`, `>`, `>=`, `IS DISTINCT FROM` and `IS NOT DISTINCT FROM`. Returns `None` +/// for any other pair of types, which the caller compares another way. +pub(crate) fn compare_float_arrays( + op: Operator, + left: &dyn Array, + right: &dyn Array, +) -> Result> { + match (left.data_type(), right.data_type()) { + (DataType::Float32, DataType::Float32) => { + arrays::(op, left.as_primitive(), right.as_primitive()).map(Some) + } + (DataType::Float64, DataType::Float64) => { + arrays::(op, left.as_primitive(), right.as_primitive()).map(Some) + } + _ => Ok(None), + } +} + +/// `left op right` for a Float32 or Float64 array and a scalar of the same type. Returns `None` for +/// any other pair of types. See [`compare_float_arrays`]. +pub(crate) fn compare_float_array_scalar( + op: Operator, + left: &dyn Array, + right: &ScalarValue, +) -> Result> { + match (left.data_type(), right) { + (DataType::Float32, ScalarValue::Float32(right)) => { + array_scalar::(op, left.as_primitive(), *right).map(Some) + } + (DataType::Float64, ScalarValue::Float64(right)) => { + array_scalar::(op, left.as_primitive(), *right).map(Some) + } + _ => Ok(None), + } +} + +fn arrays( + op: Operator, + left: &PrimitiveArray, + right: &PrimitiveArray, +) -> Result +where + T::Native: Float, +{ + if left.len() != right.len() { + return internal_err!( + "Cannot compare arrays of different lengths, got {} and {}", + left.len(), + right.len() + ); + } + let (l, r) = (left.values().as_ref(), right.values().as_ref()); + // Spark's SQL ordering as single tests, the same as `compare_floats(a, b)` followed by + // `is_eq()`, `is_ne()`, `is_lt()`, `is_le()`, `is_gt()` or `is_ge()`. IEEE 754 comparisons + // already treat `-0.0` as equal to `0.0`; the NaN terms put every NaN above the other values + // and make all NaNs equal. + let values = match op { + Operator::Eq | Operator::IsNotDistinctFrom => { + pack_pairs(l, r, |a, b| (a == b) | (a.is_nan() & b.is_nan())) + } + Operator::NotEq | Operator::IsDistinctFrom => { + pack_pairs(l, r, |a, b| (a != b) & !(a.is_nan() & b.is_nan())) + } + Operator::Lt => pack_pairs(l, r, |a, b| (a < b) | (!a.is_nan() & b.is_nan())), + Operator::LtEq => pack_pairs(l, r, |a, b| (a <= b) | b.is_nan()), + Operator::Gt => pack_pairs(l, r, |a, b| (a > b) | (a.is_nan() & !b.is_nan())), + Operator::GtEq => pack_pairs(l, r, |a, b| (a >= b) | a.is_nan()), + _ => return internal_err!("Unsupported operator for a Spark float comparison: {op}"), + }; + Ok(comparison_with_nulls( + op, + values, + left.nulls(), + right.nulls(), + )) +} + +fn array_scalar( + op: Operator, + left: &PrimitiveArray, + right: Option, +) -> Result +where + T::Native: Float, +{ + let len = left.len(); + let Some(s) = right else { + // Every comparison with a null is null, and only a null is not distinct from one. + return Ok(match op { + Operator::IsNotDistinctFrom => BooleanArray::new( + left.nulls() + .map(|nulls| !nulls.inner()) + .unwrap_or_else(|| BooleanBuffer::new_unset(len)), + None, + ), + Operator::IsDistinctFrom => BooleanArray::new( + left.nulls() + .map(|nulls| nulls.inner().clone()) + .unwrap_or_else(|| BooleanBuffer::new_set(len)), + None, + ), + Operator::Eq + | Operator::NotEq + | Operator::Lt + | Operator::LtEq + | Operator::Gt + | Operator::GtEq => BooleanArray::new_null(len), + _ => return internal_err!("Unsupported operator for a Spark float comparison: {op}"), + }); + }; + let l = left.values().as_ref(); + // With the scalar fixed, each operator reduces to one or two tests, chosen by whether the + // scalar is NaN. Every value sorts below a NaN or equals it, and a NaN sorts above any other + // value. + let values = match (op, s.is_nan()) { + (Operator::Eq | Operator::IsNotDistinctFrom, true) => pack(l, |a| a.is_nan()), + (Operator::Eq | Operator::IsNotDistinctFrom, false) => pack(l, |a| a == s), + (Operator::NotEq | Operator::IsDistinctFrom, true) => pack(l, |a| !a.is_nan()), + (Operator::NotEq | Operator::IsDistinctFrom, false) => pack(l, |a| a != s), + (Operator::Lt, true) => pack(l, |a| !a.is_nan()), + (Operator::Lt, false) => pack(l, |a| a < s), + (Operator::LtEq, true) => BooleanBuffer::new_set(len), + (Operator::LtEq, false) => pack(l, |a| a <= s), + (Operator::Gt, true) => BooleanBuffer::new_unset(len), + (Operator::Gt, false) => pack(l, |a| (a > s) | a.is_nan()), + (Operator::GtEq, true) => pack(l, |a| a.is_nan()), + (Operator::GtEq, false) => pack(l, |a| (a >= s) | a.is_nan()), + _ => return internal_err!("Unsupported operator for a Spark float comparison: {op}"), + }; + Ok(comparison_with_nulls(op, values, left.nulls(), None)) +} + +/// Applies SQL's null rules to `values`, the comparison of every slot including the null ones. +/// A null operand makes `=`, `<>`, `<`, `<=`, `>` and `>=` null. `IS NOT DISTINCT FROM`, for which +/// `values` holds equality, and `IS DISTINCT FROM`, for which it holds inequality, compare a null +/// as a value instead: equal to another null and distinct from anything else. +pub(crate) fn comparison_with_nulls( + op: Operator, + values: BooleanBuffer, + left: Option<&NullBuffer>, + right: Option<&NullBuffer>, +) -> BooleanArray { + match op { + Operator::IsNotDistinctFrom => { + let values = match (left, right) { + (None, None) => values, + (Some(nulls), None) | (None, Some(nulls)) => &values & nulls.inner(), + (Some(l), Some(r)) => { + let (l, r) = (l.inner(), r.inner()); + &(&(&values & l) & r) | &!&(l | r) + } + }; + BooleanArray::new(values, None) + } + Operator::IsDistinctFrom => { + let values = match (left, right) { + (None, None) => values, + (Some(nulls), None) | (None, Some(nulls)) => &values | &!nulls.inner(), + (Some(l), Some(r)) => { + let (l, r) = (l.inner(), r.inner()); + &(&(&values & l) & r) | &(l ^ r) + } + }; + BooleanArray::new(values, None) + } + _ => BooleanArray::new(values, NullBuffer::union(left, right)), + } +} + +/// Packs `test(value)` for each value into a bitmap, 64 values to a word. +#[inline] +fn pack(values: &[T], test: impl Fn(T) -> bool) -> BooleanBuffer { + let mut buffer = MutableBuffer::new(bit_util::ceil(values.len(), 64) * 8); + let (chunks, remainder) = values.as_chunks::<64>(); + for chunk in chunks { + let mut word = 0u64; + for (bit, &value) in chunk.iter().enumerate() { + word |= (test(value) as u64) << bit; + } + buffer.push(word); + } + if !remainder.is_empty() { + let mut word = 0u64; + for (bit, &value) in remainder.iter().enumerate() { + word |= (test(value) as u64) << bit; + } + buffer.push(word); + } + BooleanBuffer::new(buffer.into(), 0, values.len()) +} + +/// Packs `test(left[i], right[i])` for each pair into a bitmap. See [`pack`]. +#[inline] +fn pack_pairs(left: &[T], right: &[T], test: impl Fn(T, T) -> bool) -> BooleanBuffer { + debug_assert_eq!(left.len(), right.len()); + let mut buffer = MutableBuffer::new(bit_util::ceil(left.len(), 64) * 8); + let (left_chunks, left_remainder) = left.as_chunks::<64>(); + let (right_chunks, right_remainder) = right.as_chunks::<64>(); + for (l, r) in left_chunks.iter().zip(right_chunks) { + let mut word = 0u64; + for bit in 0..64 { + word |= (test(l[bit], r[bit]) as u64) << bit; + } + buffer.push(word); + } + if !left_remainder.is_empty() { + let mut word = 0u64; + for (bit, (&l, &r)) in left_remainder.iter().zip(right_remainder).enumerate() { + word |= (test(l, r) as u64) << bit; + } + buffer.push(word); + } + BooleanBuffer::new(buffer.into(), 0, left.len()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::float_semantics::{compare_floats, EDGE_VALUES}; + use arrow::array::{ArrayRef, Float32Array, Float64Array}; + use std::cmp::Ordering; + use std::sync::Arc; + + const OPERATORS: [Operator; 8] = [ + Operator::Eq, + Operator::NotEq, + Operator::Lt, + Operator::LtEq, + Operator::Gt, + Operator::GtEq, + Operator::IsDistinctFrom, + Operator::IsNotDistinctFrom, + ]; + + /// Spark's answer for `left op right` from [`compare_floats`]. + fn expected(op: Operator, left: Option, right: Option) -> Option { + let ordering = match (left, right) { + (Some(l), Some(r)) => Some(compare_floats(l, r)), + _ => None, + }; + match op { + Operator::IsNotDistinctFrom => Some(match (left, right) { + (None, None) => true, + (Some(_), Some(_)) => ordering == Some(Ordering::Equal), + _ => false, + }), + Operator::IsDistinctFrom => { + expected(Operator::IsNotDistinctFrom, left, right).map(|v| !v) + } + _ => ordering.map(|ordering| match op { + Operator::Eq => ordering.is_eq(), + Operator::NotEq => ordering.is_ne(), + Operator::Lt => ordering.is_lt(), + Operator::LtEq => ordering.is_le(), + Operator::Gt => ordering.is_gt(), + Operator::GtEq => ordering.is_ge(), + _ => unreachable!(), + }), + } + } + + /// `values` as a Float64 array, or as a Float32 array holding the same values. + fn array(values: &[Option], float32: bool) -> ArrayRef { + if float32 { + Arc::new(Float32Array::from( + values + .iter() + .map(|v| v.map(|v| v as f32)) + .collect::>(), + )) + } else { + Arc::new(Float64Array::from(values.to_vec())) + } + } + + fn scalar(value: Option, float32: bool) -> ScalarValue { + if float32 { + ScalarValue::Float32(value.map(|v| v as f32)) + } else { + ScalarValue::Float64(value) + } + } + + /// Every pair of edge values, under every operator, as two arrays and as an array and a + /// scalar, with enough rows that the comparison spans whole words and a partial one. + #[test] + fn every_pair_of_edge_values_matches_compare_floats() -> Result<()> { + // Repeat the pairs so that a run covers several 64-value words and a remainder. + let pairs: Vec<(Option, Option)> = EDGE_VALUES + .iter() + .flat_map(|&l| EDGE_VALUES.iter().map(move |&r| (l, r))) + .cycle() + .take(EDGE_VALUES.len() * EDGE_VALUES.len() * 3 + 7) + .collect(); + let left: Vec<_> = pairs.iter().map(|p| p.0).collect(); + let right: Vec<_> = pairs.iter().map(|p| p.1).collect(); + for float32 in [false, true] { + // Float32 holds every edge value exactly, NaNs included, as the same class of value. + let (l, r) = (array(&left, float32), array(&right, float32)); + for op in OPERATORS { + let actual = compare_float_arrays(op, l.as_ref(), r.as_ref())?.unwrap(); + let want: Vec<_> = pairs.iter().map(|&(a, b)| expected(op, a, b)).collect(); + assert_eq!( + actual.iter().collect::>(), + want, + "a {op} b, f32={float32}" + ); + for &value in &EDGE_VALUES { + let s = scalar(value, float32); + let actual = compare_float_array_scalar(op, l.as_ref(), &s)?.unwrap(); + let want: Vec<_> = left.iter().map(|&a| expected(op, a, value)).collect(); + assert_eq!( + actual.iter().collect::>(), + want, + "a {op} {value:?}, f32={float32}" + ); + } + } + } + Ok(()) + } + + /// A sliced array starts its values and its null bitmap at an offset. + #[test] + fn sliced_inputs() -> Result<()> { + let values: Vec> = (0..200) + .map(|i| EDGE_VALUES[(i * 7) % EDGE_VALUES.len()]) + .collect(); + let others: Vec> = (0..200) + .map(|i| EDGE_VALUES[(i * 3 + 1) % EDGE_VALUES.len()]) + .collect(); + let (l, r) = (array(&values, false), array(&others, false)); + for (offset, len) in [(1, 150), (63, 70), (64, 64), (130, 5), (3, 0)] { + let (ls, rs) = (l.slice(offset, len), r.slice(offset + 2, len)); + for op in OPERATORS { + let actual = compare_float_arrays(op, ls.as_ref(), rs.as_ref())?.unwrap(); + let want: Vec<_> = (0..len) + .map(|i| expected(op, values[offset + i], others[offset + 2 + i])) + .collect(); + assert_eq!(actual.iter().collect::>(), want, "{op} at {offset}"); + let actual = + compare_float_array_scalar(op, ls.as_ref(), &ScalarValue::Float64(None))? + .unwrap(); + let want: Vec<_> = (0..len) + .map(|i| expected(op, values[offset + i], None)) + .collect(); + assert_eq!( + actual.iter().collect::>(), + want, + "{op} null at {offset}" + ); + } + } + Ok(()) + } + + #[test] + fn other_types_are_left_to_the_caller() -> Result<()> { + let floats: ArrayRef = Arc::new(Float64Array::from(vec![1.0])); + let narrow: ArrayRef = Arc::new(Float32Array::from(vec![1.0])); + assert!(compare_float_arrays(Operator::Lt, floats.as_ref(), narrow.as_ref())?.is_none()); + assert!(compare_float_array_scalar( + Operator::Lt, + floats.as_ref(), + &ScalarValue::Float32(Some(1.0)) + )? + .is_none()); + assert!( + compare_float_array_scalar(Operator::Lt, floats.as_ref(), &ScalarValue::Null)? + .is_none() + ); + assert!(compare_float_arrays(Operator::Plus, floats.as_ref(), floats.as_ref()).is_err()); + let different_lengths: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0])); + assert!( + compare_float_arrays(Operator::Lt, floats.as_ref(), different_lengths.as_ref()) + .is_err() + ); + Ok(()) + } +} diff --git a/native/spark-expr/src/float_semantics/mod.rs b/native/spark-expr/src/float_semantics/mod.rs index 60f8614c285..997f8b24648 100644 --- a/native/spark-expr/src/float_semantics/mod.rs +++ b/native/spark-expr/src/float_semantics/mod.rs @@ -37,21 +37,27 @@ //! [`hash_input`]. //! //! A native expression must follow the rule of the Spark function it replaces, so build it from -//! these helpers rather than a local copy. Where an Arrow kernel sorts, row-encodes or hashes the -//! values, normalize them first: once `-0.0` is folded and NaN canonicalized, Arrow's total order -//! agrees with `compareDoubles`. Where Comet compares values itself, or has to return the original -//! bits as `array_min` does, use [`compare_floats`], [`float_lt`], [`float_gt`], -//! [`spark_comparator`] or [`spark_equality`]. Non-canonical NaNs are not a corner case: on x86-64 -//! every NaN that arithmetic produces at run time, such as `sqrt(-1)`, has the sign bit set. +//! these helpers rather than a local copy. Where an Arrow kernel sorts, row-encodes, hashes or +//! compares the values, normalize them first: once `-0.0` is folded and NaN canonicalized, Arrow's +//! total order agrees with `compareDoubles`. Where Comet compares values itself, or has to return +//! the original bits as `array_min` does, use [`compare_floats`], [`float_lt`], [`float_gt`], +//! [`spark_comparator`] or [`spark_equality`]. The comparison operators compare in place too, with +//! the kernels in `kernels` for flat floats and with [`spark_comparator`] and [`spark_equality`] +//! for lists and structs. Non-canonical NaNs are not a corner case: on x86-64 every NaN that +//! arithmetic produces at run time, such as `sqrt(-1)`, has the sign bit set. mod compare; +mod kernels; mod normalize; pub use compare::{spark_comparator, spark_equality}; +pub(crate) use compare::{spark_comparator_ignoring_nulls, spark_equality_ignoring_nulls}; +pub(crate) use kernels::{compare_float_array_scalar, compare_float_arrays, comparison_with_nulls}; pub use normalize::{ - has_float_leaf, normalize_floats, normalize_nested_floats, NormalizeNaNAndZero, - NormalizeNestedFloats, + has_float_leaf, normalize_comparison_operand, normalize_floats, normalize_nested_floats, + NormalizeNaNAndZero, NormalizeNestedFloats, }; +pub(crate) use normalize::{is_nested_with_float_leaf, normalize_float_scalar}; use num::Float; use std::cmp::Ordering; diff --git a/native/spark-expr/src/float_semantics/normalize.rs b/native/spark-expr/src/float_semantics/normalize.rs index b518e826f04..33b1c27b5ff 100644 --- a/native/spark-expr/src/float_semantics/normalize.rs +++ b/native/spark-expr/src/float_semantics/normalize.rs @@ -24,6 +24,7 @@ use arrow::datatypes::{DataType, Float32Type, Float64Type, Schema}; use arrow::record_batch::RecordBatch; use datafusion::common::{Result, ScalarValue}; use datafusion::logical_expr::ColumnarValue; +use datafusion::physical_expr::expressions::Literal; use datafusion::physical_expr::PhysicalExpr; use std::fmt::{Display, Formatter}; use std::hash::{Hash, Hasher}; @@ -92,15 +93,17 @@ impl PhysicalExpr for NormalizeNaNAndZero { } fn evaluate(&self, batch: &RecordBatch) -> Result { - let cv = self.child.evaluate(batch)?; - let array = cv.into_array(batch.num_rows())?; - match &self.data_type { - DataType::Float32 | DataType::Float64 => { - Ok(ColumnarValue::Array(normalize_floats(&array))) - } + DataType::Float32 | DataType::Float64 => {} dt => panic!("Unexpected data type {dt:?}"), } + // A scalar stays a scalar, so a constant operand is not expanded into a column. + match self.child.evaluate(batch)? { + ColumnarValue::Array(array) => Ok(ColumnarValue::Array(normalize_floats(&array))), + ColumnarValue::Scalar(value) => { + Ok(ColumnarValue::Scalar(normalize_float_scalar(value))) + } + } } fn children(&self) -> Vec<&Arc> { @@ -124,7 +127,8 @@ impl Display for NormalizeNaNAndZero { } } -/// Normalizes nested IN operands, preserving constants for static membership lookup. +/// Applies [`normalize_nested_floats`] to a list or struct operand of `IN` or a comparison. A +/// constant operand stays a scalar, which keeps static membership lookup available. #[derive(Debug, Eq)] pub struct NormalizeNestedFloats { child: Arc, @@ -132,19 +136,13 @@ pub struct NormalizeNestedFloats { impl NormalizeNestedFloats { /// Wrap nested floating-point operands only; scalar floats keep their existing semantics. + /// An operand that is already wrapped is returned as is. pub fn wrap_if_needed( child: Arc, schema: &Schema, ) -> Result> { let dt = child.data_type(schema)?; - if matches!( - dt, - DataType::List(_) - | DataType::LargeList(_) - | DataType::FixedSizeList(_, _) - | DataType::Struct(_) - ) && has_float_leaf(&dt) - { + if is_nested_with_float_leaf(&dt) && child.downcast_ref::().is_none() { Ok(Arc::new(Self { child })) } else { Ok(child) @@ -214,6 +212,32 @@ impl PhysicalExpr for NormalizeNestedFloats { } } +/// Normalizes an operand of a comparison, so that Arrow's comparison kernels order it the way +/// Spark's SQL ordering does. Arrow orders floats by IEEE 754 total order, which agrees with +/// Spark's once `-0.0` is folded and every NaN canonicalized. +/// +/// A Float32 or Float64 operand is wrapped in [`NormalizeNaNAndZero`], and a list or struct with +/// a float leaf in [`NormalizeNestedFloats`]. A literal is normalized now instead, so that it stays +/// a literal. Other operands, and operands that are already wrapped, are returned as is. +pub fn normalize_comparison_operand( + operand: Arc, + schema: &Schema, +) -> Result> { + if let Some(literal) = operand.downcast_ref::() { + let value = literal.value(); + let normalized = match value.data_type() { + DataType::Float32 | DataType::Float64 => normalize_float_scalar(value.clone()), + dt if is_nested_with_float_leaf(&dt) => { + ScalarValue::try_from_array(&normalize_nested_floats(&value.to_array()?), 0)? + } + _ => return Ok(operand), + }; + return Ok(Arc::new(Literal::new(normalized))); + } + let operand = NormalizeNaNAndZero::wrap_if_needed(operand, schema)?; + NormalizeNestedFloats::wrap_if_needed(operand, schema) +} + /// Applies [`normalize_float`] to a Float32 or Float64 array. Any other array is returned as is, /// including a nested one: [`normalize_nested_floats`] reaches floats inside lists and structs. pub fn normalize_floats(array: &ArrayRef) -> ArrayRef { @@ -232,6 +256,15 @@ pub fn normalize_floats(array: &ArrayRef) -> ArrayRef { } } +/// Applies [`normalize_float`] to a Float32 or Float64 scalar. Other scalars are returned as is. +pub(crate) fn normalize_float_scalar(value: ScalarValue) -> ScalarValue { + match value { + ScalarValue::Float32(v) => ScalarValue::Float32(v.map(normalize_float)), + ScalarValue::Float64(v) => ScalarValue::Float64(v.map(normalize_float)), + other => other, + } +} + /// Whether a Float32 or Float64 field occurs at any depth of a list or struct type. The keys and /// values of a map are not searched. pub fn has_float_leaf(dt: &DataType) -> bool { @@ -245,6 +278,11 @@ pub fn has_float_leaf(dt: &DataType) -> bool { } } +/// Whether `dt` is a list or struct with a Float32 or Float64 field at any depth. +pub(crate) fn is_nested_with_float_leaf(dt: &DataType) -> bool { + dt.is_nested() && has_float_leaf(dt) +} + /// Recursively rebuilds nested arrays with `-0.0` normalized to `0.0` and NaN canonicalized /// in any Float32/Float64 leaves. /// @@ -526,6 +564,120 @@ mod tests { assert!(Arc::ptr_eq(&normalize_floats(&nested), &nested)); } + #[test] + fn test_normalize_nan_and_zero_keeps_scalars() -> Result<()> { + use datafusion::physical_expr::expressions::Literal; + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + )?; + for (value, expected) in [ + ( + ScalarValue::Float64(Some(-0.0)), + ScalarValue::Float64(Some(0.0)), + ), + ( + ScalarValue::Float64(Some(f64::from_bits(0xfff8_0000_0000_0000))), + ScalarValue::Float64(Some(f64::NAN)), + ), + ( + ScalarValue::Float32(Some(-0.0)), + ScalarValue::Float32(Some(0.0)), + ), + (ScalarValue::Float64(None), ScalarValue::Float64(None)), + ] { + let expr = + NormalizeNaNAndZero::new(value.data_type(), Arc::new(Literal::new(value.clone()))); + // Compare the bits, since every NaN prints the same. + let bits = |v: &ScalarValue| match v { + ScalarValue::Float64(v) => v.map(f64::to_bits), + ScalarValue::Float32(v) => v.map(|v| u64::from(v.to_bits())), + other => panic!("unexpected {other:?}"), + }; + match expr.evaluate(&batch)? { + ColumnarValue::Scalar(actual) => { + assert_eq!(actual.data_type(), expected.data_type()); + assert_eq!(bits(&actual), bits(&expected), "{value:?}"); + } + ColumnarValue::Array(_) => panic!("{value:?} was expanded into an array"), + } + } + Ok(()) + } + + #[test] + fn test_normalize_comparison_operand() -> Result<()> { + use datafusion::physical_expr::expressions::{Column, Literal}; + let list = DataType::List(Arc::new(Field::new("item", DataType::Float64, true))); + let schema = Schema::new(vec![ + Field::new("f64", DataType::Float64, true), + Field::new("list", list, true), + Field::new("i32", DataType::Int32, true), + ]); + let column = |i: usize| -> Arc { + Arc::new(Column::new(schema.field(i).name(), i)) + }; + + // Columns are wrapped once, and only when they hold floats. + let wrapped = normalize_comparison_operand(column(0), &schema)?; + assert!(wrapped.downcast_ref::().is_some()); + let again = normalize_comparison_operand(Arc::clone(&wrapped), &schema)?; + assert!(Arc::ptr_eq(&again, &wrapped)); + let wrapped = normalize_comparison_operand(column(1), &schema)?; + assert!(wrapped.downcast_ref::().is_some()); + let again = normalize_comparison_operand(Arc::clone(&wrapped), &schema)?; + assert!(Arc::ptr_eq(&again, &wrapped)); + let int = column(2); + assert!(Arc::ptr_eq( + &normalize_comparison_operand(Arc::clone(&int), &schema)?, + &int + )); + + // Literals are normalized now and stay literals. + let literal = + |value: ScalarValue| -> Arc { Arc::new(Literal::new(value)) }; + let value = |expr: Arc| -> ScalarValue { + expr.downcast_ref::() + .expect("a literal operand stays a literal") + .value() + .clone() + }; + let bits = |expr: Arc| match value(expr) { + ScalarValue::Float64(Some(v)) => v.to_bits(), + ScalarValue::Float32(Some(v)) => u64::from(v.to_bits()), + other => panic!("unexpected {other:?}"), + }; + let folded = + normalize_comparison_operand(literal(ScalarValue::Float64(Some(-0.0))), &schema)?; + assert_eq!(bits(folded), 0.0f64.to_bits()); + let folded = normalize_comparison_operand( + literal(ScalarValue::Float32(Some(f32::from_bits(0xffc0_0000)))), + &schema, + )?; + assert_eq!(bits(folded), u64::from(f32::NAN.to_bits())); + let nested = ScalarValue::List(Arc::new( + ListArray::from_iter_primitive::([Some(vec![ + Some(-0.0), + Some(f64::from_bits(0xfff8_0000_0000_0000)), + ])]), + )); + let folded = value(normalize_comparison_operand(literal(nested), &schema)?); + let ScalarValue::List(folded) = folded else { + panic!("expected a list literal, got {folded:?}"); + }; + let folded = folded.value(0); + let folded = folded.as_primitive::(); + assert_eq!(folded.value(0).to_bits(), 0.0f64.to_bits()); + assert_eq!(folded.value(1).to_bits(), f64::NAN.to_bits()); + let int = literal(ScalarValue::Int32(Some(1))); + assert!(Arc::ptr_eq( + &normalize_comparison_operand(Arc::clone(&int), &schema)?, + &int + )); + Ok(()) + } + #[test] fn test_wrap_float_keys_once() -> Result<()> { use datafusion::physical_expr::expressions::Column; diff --git a/native/spark-expr/src/json_funcs/to_json.rs b/native/spark-expr/src/json_funcs/to_json.rs index caf87514fad..94f29e3f458 100644 --- a/native/spark-expr/src/json_funcs/to_json.rs +++ b/native/spark-expr/src/json_funcs/to_json.rs @@ -139,7 +139,7 @@ fn array_to_json_string( spark_cast( ColumnarValue::Array(Arc::clone(arr)), &DataType::Utf8, - &SparkCastOptions::new(EvalMode::Legacy, timezone, false), + &SparkCastOptions::new(EvalMode::Legacy, timezone), )? .into_array(arr.len()) } diff --git a/native/spark-expr/src/jvm_udf/mod.rs b/native/spark-expr/src/jvm_udf/mod.rs index 232926744d0..6355b1a9da6 100644 --- a/native/spark-expr/src/jvm_udf/mod.rs +++ b/native/spark-expr/src/jvm_udf/mod.rs @@ -26,6 +26,7 @@ use arrow::record_batch::RecordBatch; use datafusion::common::Result as DFResult; use datafusion::logical_expr::ColumnarValue; +use datafusion::physical_expr::expressions::Literal; use datafusion::physical_expr::PhysicalExpr; use datafusion_comet_common::{decode_string_arrays, zero_offsets}; @@ -39,6 +40,12 @@ use jni::objects::{Global, JObject, JValue}; pub struct JvmScalarUdfExpr { class_name: String, args: Vec>, + /// The length-1 array sent for each argument that is a literal. Built once here, because + /// building it for every batch copies the value twice: `Literal::evaluate` clones it and + /// `to_array_of_size` copies the clone. The codegen dispatcher's serialized expression is + /// such a literal, several KB for a Scala UDF. `None` for any other argument, and for a + /// literal whose array cannot be built, which `evaluate` then reports. + literal_arrays: Vec>, return_type: DataType, return_nullable: bool, /// Captured at `createPlan` time and threaded here by the planner. Passed through the @@ -54,9 +61,14 @@ pub struct JvmScalarUdfExpr { /// duration of the call. `None` when no driving Spark task is available (unit tests, direct /// native driver runs); the bridge then installs nothing. class_loader: Option>>>, + /// Index of the partition this native plan computes. See `CometUDF.evaluate`. + partition: i32, + /// Id of this native plan. See `CometUDF.evaluate`. + exec_context_id: i64, } impl JvmScalarUdfExpr { + #[allow(clippy::too_many_arguments)] pub fn new( class_name: String, args: Vec>, @@ -64,18 +76,30 @@ impl JvmScalarUdfExpr { return_nullable: bool, task_context: Option>>>, class_loader: Option>>>, + partition: i32, + exec_context_id: i64, ) -> Self { debug_assert!( !class_name.is_empty(), "JvmScalarUdfExpr requires a non-empty class name" ); + let literal_arrays = args + .iter() + .map(|arg| { + arg.downcast_ref::() + .and_then(|literal| literal.value().to_array_of_size(1).ok()) + }) + .collect(); Self { class_name, args, + literal_arrays, return_type, return_nullable, task_context, class_loader, + partition, + exec_context_id, } } } @@ -134,9 +158,13 @@ impl PhysicalExpr for JvmScalarUdfExpr { let arrays: Vec = self .args .iter() - .map(|e| match e.evaluate(batch)? { - ColumnarValue::Array(a) => Ok(a), - ColumnarValue::Scalar(s) => s.to_array_of_size(1), + .zip(&self.literal_arrays) + .map(|(e, literal_array)| match literal_array { + Some(a) => Ok(Arc::clone(a)), + None => match e.evaluate(batch)? { + ColumnarValue::Array(a) => Ok(a), + ColumnarValue::Scalar(s) => s.to_array_of_size(1), + }, }) .collect::>()?; @@ -234,6 +262,8 @@ impl PhysicalExpr for JvmScalarUdfExpr { JValue::Long(out_arr_ptr).as_jni(), JValue::Long(out_sch_ptr).as_jni(), JValue::Int(batch.num_rows() as i32).as_jni(), + JValue::Int(self.partition).as_jni(), + JValue::Long(self.exec_context_id).as_jni(), JValue::Object(task_context_ref).as_jni(), JValue::Object(class_loader_ref).as_jni(), ], @@ -276,6 +306,53 @@ impl PhysicalExpr for JvmScalarUdfExpr { self.return_nullable, self.task_context.clone(), self.class_loader.clone(), + self.partition, + self.exec_context_id, ))) } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Array, BinaryArray}; + use datafusion::common::ScalarValue; + use datafusion::physical_expr::expressions::Column; + + #[test] + fn builds_literal_arguments_once() { + let bytes = vec![7u8; 4096]; + let udf = JvmScalarUdfExpr::new( + "org.example.Udf".to_string(), + vec![ + Arc::new(Literal::new(ScalarValue::Binary(Some(bytes.clone())))), + Arc::new(Column::new("a", 0)), + ], + DataType::Int64, + true, + None, + None, + 0, + 0, + ); + let literal = udf.literal_arrays[0] + .as_ref() + .expect("a literal argument's array is built up front"); + let literal = literal.as_any().downcast_ref::().unwrap(); + assert_eq!(literal.len(), 1); + assert_eq!(literal.value(0), bytes.as_slice()); + assert!(udf.literal_arrays[1].is_none()); + + // New children rebuild the arrays, so a position that no longer holds a literal does not + // keep sending the old literal's array. + let swapped = Arc::new(udf) + .with_new_children(vec![ + Arc::new(Column::new("a", 0)), + Arc::new(Literal::new(ScalarValue::Binary(Some(bytes)))), + ]) + .unwrap(); + let swapped = swapped.downcast_ref::().unwrap(); + assert!(swapped.literal_arrays[0].is_none()); + assert!(swapped.literal_arrays[1].is_some()); + } +} diff --git a/native/spark-expr/src/kernels/temporal.rs b/native/spark-expr/src/kernels/temporal.rs index a7332d2099d..ce6ad9de5a9 100644 --- a/native/spark-expr/src/kernels/temporal.rs +++ b/native/spark-expr/src/kernels/temporal.rs @@ -18,7 +18,8 @@ //! temporal kernels use chrono::{ - DateTime, Datelike, Duration, LocalResult, NaiveDate, NaiveDateTime, TimeZone, Timelike, Utc, + DateTime, Datelike, Duration, LocalResult, NaiveDate, NaiveDateTime, Offset, TimeZone, + Timelike, Utc, }; use std::sync::Arc; @@ -32,11 +33,54 @@ use arrow::array::{ }; use arrow::{ array::*, - datatypes::{DataType, TimeUnit}, + datatypes::{DataType, Field, TimeUnit}, }; +use datafusion::{ + common::{config::ConfigOptions, ScalarValue}, + logical_expr::{ColumnarValue, ScalarFunctionArgs}, +}; +use datafusion_functions::datetime; use crate::SparkError; +/// Invoke DataFusion's physical `date_trunc` implementation with a scalar granularity. +/// +/// Spark syntax normalization and compatibility fallback decisions deliberately live outside +/// this helper so the upstream execution boundary stays obvious. +fn datafusion_date_trunc( + array: ArrayRef, + granularity: &'static str, +) -> Result { + let data_type = array.data_type().clone(); + let number_rows = array.len(); + let args = vec![ + ColumnarValue::Scalar(ScalarValue::Utf8(Some(granularity.to_string()))), + ColumnarValue::Array(array), + ]; + let arg_fields = args + .iter() + .enumerate() + .map(|(index, value)| { + Arc::new(Field::new( + format!("date_trunc_arg_{index}"), + value.data_type(), + true, + )) + }) + .collect(); + + datetime::date_trunc() + .invoke_with_args(ScalarFunctionArgs { + args, + arg_fields, + number_rows, + return_field: Arc::new(Field::new("date_trunc", data_type, true)), + config_options: Arc::new(ConfigOptions::default()), + }) + .and_then(|value| value.to_array(number_rows)) + .map_err(|error| SparkError::Internal(error.to_string())) +} + // Copied from arrow_arith/temporal.rs macro_rules! return_compute_error_with { ($msg:expr, $param:expr) => { @@ -47,6 +91,7 @@ macro_rules! return_compute_error_with { // The number of days between the beginning of the proleptic gregorian calendar (0001-01-01) // and the beginning of the Unix Epoch (1970-01-01) const DAYS_TO_UNIX_EPOCH: i32 = 719_163; +const MICROS_PER_MINUTE: i64 = 60_000_000; // Optimized date truncation functions that work directly with days since epoch // These avoid the overhead of converting to/from NaiveDateTime @@ -132,30 +177,6 @@ where Ok(builder.finish().with_timezone(tz_str)) } -fn as_timestamp_tz_with_op_single( - value: Option, - builder: &mut PrimitiveBuilder, - tz: &Tz, - op: F, -) -> Result<(), SparkError> -where - F: Fn(DateTime) -> i64, - i64: From, -{ - match value { - Some(value) => match as_datetime_with_timezone::(value.into(), *tz) { - Some(time) => builder.append_value(op(time)), - _ => { - return Err(SparkError::Internal( - "Unable to read value as datetime".to_string(), - )); - } - }, - None => builder.append_null(), - } - Ok(()) -} - // Apply the Tz to the Naive Date Time, convert to UTC, and return as microseconds in Unix epoch. // After truncation the carried UTC offset may be wrong if the truncated time falls in a different // DST period than the original (e.g., truncating a December/PST timestamp to QUARTER yields @@ -353,19 +374,100 @@ fn date_trunc_fn_for_format(format: &str) -> Result { }) } -/// Optimized date truncation for Date32 arrays -/// Works directly with days since epoch instead of converting to/from NaiveDateTime -fn date_trunc_date32(array: &Date32Array, format: String) -> Result { - // Select the truncation function based on format - let trunc_fn = date_trunc_fn_for_format(&format)?; +const MICROS_PER_DAY: i64 = 86_400_000_000; - // Apply truncation to each element - let result: Date32Array = array - .iter() - .map(|opt_days| opt_days.and_then(trunc_fn)) - .collect(); +#[inline] +fn fits_timestamp_nanosecond(micros: i64) -> bool { + micros.checked_mul(1_000).is_some() +} - Ok(result) +/// DataFusion's coarse truncation first converts the input to nanoseconds. Although an input near +/// the lower TimestampNanosecond bound can itself be represented, truncating it may move the +/// result before that bound: for example, `1677-09-22` truncated to YEAR becomes `1677-01-01`. +/// The result can move backward by 365 days (366 when starting from December 31 in a leap year), +/// and timezone gap handling can shift it by a few more hours. Round that worst case up to 370 +/// days so both DataFusion's input and coarse-truncation result remain representable. The +/// effective microsecond interval is approximately `1678-09-26T00:12:43.145225Z` through +/// `2262-04-11T23:47:16.854775Z`; because Date32 values are UTC midnight, its first upstream date +/// is 1678-09-27 and its last is 2262-04-11. +#[inline] +fn fits_datafusion_coarse_trunc_range(micros: i64) -> bool { + const LOWER_NANOSECOND_MICROS: i64 = i64::MIN / 1_000; + const COARSE_TRUNC_MARGIN_MICROS: i64 = 370 * MICROS_PER_DAY; + + fits_timestamp_nanosecond(micros) + && micros >= LOWER_NANOSECOND_MICROS + COARSE_TRUNC_MARGIN_MICROS +} + +// Integer proleptic Gregorian calendar conversion, adapted from DataFusion 55.1.0's +// datetime/date_trunc.rs (Howard Hinnant's algorithms): +// https://howardhinnant.github.io/date_algorithms.html +// These helpers only receive days derived from i64 microseconds (about +/-107 million), +// so the calendar intermediates fit in i64 even outside chrono's year range. +fn civil_from_days(days: i64) -> (i64, i64, i64) { + let z = days + 719_468; + let era = z.div_euclid(146_097); + let day_of_era = z.rem_euclid(146_097); + let year_of_era = + (day_of_era - day_of_era / 1460 + day_of_era / 36524 - day_of_era / 146_096) / 365; + let day_of_year = day_of_era - (365 * year_of_era + year_of_era / 4 - year_of_era / 100); + let month_index = (5 * day_of_year + 2) / 153; + let day = day_of_year - (153 * month_index + 2) / 5 + 1; + let month = if month_index < 10 { + month_index + 3 + } else { + month_index - 9 + }; + let year = year_of_era + era * 400 + i64::from(month <= 2); + (year, month, day) +} + +fn days_from_civil(year: i64, month: i64, day: i64) -> i64 { + let year = year - i64::from(month <= 2); + let era = year.div_euclid(400); + let year_of_era = year.rem_euclid(400); + let month_index = if month > 2 { month - 3 } else { month + 9 }; + let day_of_year = (153 * month_index + 2) / 5 + day - 1; + let day_of_era = year_of_era * 365 + year_of_era / 4 - year_of_era / 100 + day_of_year; + era * 146_097 + day_of_era - 719_468 +} + +/// Truncate UTC/NTZ coarse units directly in microseconds, avoiding both DataFusion's +/// intermediate nanosecond conversion and chrono's year limit. A representable input can +/// truncate below i64::MIN; report that overflow rather than returning NULL or wrapping. +fn timestamp_trunc_coarse_micros(micros: i64, granularity: &str) -> Result { + let days = micros.div_euclid(MICROS_PER_DAY); + let truncated_days = match granularity { + // The epoch is Thursday, three days after Monday. + "week" => days - (days + 3).rem_euclid(7), + "month" => { + let (_, _, day) = civil_from_days(days); + days - (day - 1) + } + "quarter" => { + let (year, month, _) = civil_from_days(days); + days_from_civil(year, 1 + 3 * ((month - 1) / 3), 1) + } + "year" => { + let (year, _, _) = civil_from_days(days); + days_from_civil(year, 1, 1) + } + _ => unreachable!("integer fallback only handles normalized coarse units"), + }; + truncated_days.checked_mul(MICROS_PER_DAY).ok_or_else(|| { + SparkError::Internal(format!( + "long overflow: Timestamp {micros} out of range after date_trunc({granularity})" + )) + }) +} + +/// Truncate Date32 values directly in days since the epoch. +/// +/// Routing Date32 through DataFusion's timestamp kernel requires two casts and temporary arrays, +/// which is materially slower than this single-pass implementation. +fn date_trunc_date32(array: &Date32Array, format: String) -> Result { + let trunc_fn = date_trunc_fn_for_format(&format)?; + Ok(array.iter().map(|value| value.and_then(trunc_fn)).collect()) } /// @@ -540,15 +642,51 @@ where /// timezone or no timezone. The array may be a dictionary array. /// /// format is a scalar string specifying the format to apply to the timestamp value. +/// +/// wrap_second_millisecond_overflow selects how SECOND and MILLISECOND truncation treats a +/// result below the smallest timestamp: Spark before 4.2 wraps, 4.2 and later raise. pub(crate) fn timestamp_trunc_dyn( array: &dyn Array, format: String, + wrap_second_millisecond_overflow: bool, ) -> Result { match array.data_type().clone() { DataType::Dictionary(_, _) => { downcast_dictionary_array!( array => { - let truncated_values = timestamp_trunc_dyn(array.values(), format)?; + // Safe, low-cardinality truncation can operate only on distinct values, + // rather than scanning every key just to discover unused values. + if !timestamp_trunc_dictionary_needs_mask( + array.values(), + array.len(), + &format, + wrap_second_millisecond_overflow, + )? { + let values = timestamp_trunc_dyn( + array.values(), + format, + wrap_second_millisecond_overflow, + )?; + return Ok(Arc::new(array.with_values(values))); + } + // Dictionary values can outlive the rows that reference them (e.g. after + // filtering). Unused entries, including entries hidden by NULL keys, + // must not raise errors in the fallible timestamp kernel. + let mut unused = vec![true; array.values().len()]; + for key in array.keys_iter().flatten() { + unused[key] = false; + } + let values = if unused.iter().any(|unused| *unused) { + arrow::compute::nullif(array.values(), &BooleanArray::from(unused)) + .map_err(|error| SparkError::Internal(error.to_string()))? + } else { + Arc::clone(array.values()) + }; + let truncated_values = timestamp_trunc_dyn( + values.as_ref(), + format, + wrap_second_millisecond_overflow, + )?; Ok(Arc::new(array.with_values(truncated_values))) } dt => return_compute_error_with!("timestamp_trunc does not support", dt), @@ -557,7 +695,7 @@ pub(crate) fn timestamp_trunc_dyn( _ => { downcast_temporal_array!( array => { - timestamp_trunc(array, format) + timestamp_trunc(array, format, wrap_second_millisecond_overflow) .map(|a| Arc::new(a) as ArrayRef) } dt => return_compute_error_with!("timestamp_trunc does not support", dt), @@ -566,6 +704,52 @@ pub(crate) fn timestamp_trunc_dyn( } } +/// Keep key masking for fallible values and for coarse units with many distinct values, where +/// masking can avoid expensive calendar work on unused entries. Fine arithmetic processes the +/// physical values regardless of validity, so infallible fine units never need a key scan. +/// MICROSECOND is always infallible. SECOND and MILLISECOND ignore the timezone and, when they do +/// not wrap, fail only near the lower bound, so only such a value needs the key scan. +fn timestamp_trunc_dictionary_needs_mask( + values: &dyn Array, + keys_len: usize, + format: &str, + wrap_second_millisecond_overflow: bool, +) -> Result { + let Some(values) = values.as_any().downcast_ref::() else { + return Ok(true); + }; + // Truncation moves backwards by at most a year. Above this conservative lower-bound margin, + // UTC/NTZ truncation cannot underflow i64 and all values can be evaluated without a key scan. + const LOWER_SAFE_MICROS: i64 = i64::MIN + 370 * MICROS_PER_DAY; + let has_value_near_lower_bound = || { + values + .iter() + .flatten() + .any(|micros| micros < LOWER_SAFE_MICROS) + }; + let granularity = normalize_timestamp_trunc_format(format)?; + if granularity == "microsecond" + || (wrap_second_millisecond_overflow && matches!(granularity, "millisecond" | "second")) + { + return Ok(false); + } + if matches!(granularity, "millisecond" | "second") { + // These units ignore the timezone, so a non-UTC value needs no masking of its own. + return Ok(has_value_near_lower_bound()); + } + if values.timezone().is_some_and(|tz| !is_utc_timezone(tz)) { + // Unused non-UTC values can also hit chrono boundary panics, not just Result errors. + return Ok(true); + } + if matches!(granularity, "week" | "month" | "quarter" | "year") && values.len() > keys_len / 64 + { + // Require substantial repetition before speculatively computing unused calendar values. + // Dense NULLs/high cardinality otherwise benefit from masking before truncation. + return Ok(true); + } + Ok(has_value_near_lower_bound()) +} + /// Convert microseconds since epoch to NaiveDateTime #[inline] fn micros_to_naive(micros: i64) -> Option { @@ -584,6 +768,25 @@ type NtzTruncFn = fn(NaiveDateTime) -> Option; /// Truncates a `DateTime`, returning `None` if the result is out of range. type TzTruncFn = fn(DateTime) -> Option>; +/// The Spark `date_trunc` spellings and their canonical DataFusion granularities. +const TIMESTAMP_TRUNC_ALIASES: [(&str, &str); 15] = [ + ("YEAR", "year"), + ("YYYY", "year"), + ("YY", "year"), + ("QUARTER", "quarter"), + ("MONTH", "month"), + ("MON", "month"), + ("MM", "month"), + ("WEEK", "week"), + ("DAY", "day"), + ("DD", "day"), + ("HOUR", "hour"), + ("MINUTE", "minute"), + ("SECOND", "second"), + ("MILLISECOND", "millisecond"), + ("MICROSECOND", "microsecond"), +]; + /// The `timestamp_trunc` formats Spark accepts for the NTZ path, and the truncation each one /// selects. All entries are ASCII, so `eq_ignore_ascii_case` on the raw input matches Spark /// without allocating. @@ -653,6 +856,19 @@ fn tz_trunc_fn_for_format(format: &str) -> Result { }) } +/// Normalize Spark `date_trunc` aliases without accepting additional DataFusion spellings. +fn normalize_timestamp_trunc_format(format: &str) -> Result<&'static str, SparkError> { + TIMESTAMP_TRUNC_ALIASES + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case(format)) + .map(|(_, granularity)| *granularity) + .ok_or_else(|| { + SparkError::Internal(format!( + "Unsupported format: {format:?} for function 'timestamp_trunc'" + )) + }) +} + /// Truncate a TimestampNTZ array without any timezone conversion. /// NTZ values are timezone-independent; we treat the raw microseconds as a naive datetime. fn timestamp_trunc_ntz( @@ -680,50 +896,347 @@ where Ok(result) } -/// Truncate a single NTZ value and append to builder -fn timestamp_trunc_ntz_single( - value: Option, - builder: &mut PrimitiveBuilder, - op: F, -) -> Result<(), SparkError> -where - F: Fn(NaiveDateTime) -> Option, -{ - match value { - Some(micros) => match micros_to_naive(micros).and_then(op) { - Some(truncated) => builder.append_value(naive_to_micros(truncated)), - None => { - return Err(SparkError::Internal( - "Unable to truncate NTZ timestamp".to_string(), - )) +/// The zone-aware fallback for HOUR/DAY outside DataFusion 55.1's internal TimestampNanosecond +/// range, for literal formats and format columns alike. UTC/NTZ coarse units use integer calendar +/// arithmetic instead. Like every native local-time conversion, it stops applying DST after +/// chrono-tz's last transition, around 2099, while Spark keeps applying the zone's rules (#6816). +fn timestamp_trunc_legacy( + array: &TimestampMicrosecondArray, + format: &str, +) -> Result { + let builder = TimestampMicrosecondBuilder::with_capacity(array.len()); + let iter = ArrayIter::new(array); + match array.data_type() { + DataType::Timestamp(TimeUnit::Microsecond, None) => { + timestamp_trunc_ntz(array, format.to_string()) + } + DataType::Timestamp(TimeUnit::Microsecond, Some(tz)) => { + let trunc_fn = tz_trunc_fn_for_format(format)?; + as_timestamp_tz_with_op::<&TimestampMicrosecondArray, TimestampMicrosecondType, _>( + iter, + builder, + tz, + |dt| as_micros_from_unix_epoch_utc(trunc_fn(dt)), + ) + } + dt => return_compute_error_with!( + "Unsupported input type '{:?}' for function 'timestamp_trunc'", + dt + ), + } +} + +fn datafusion_timestamp_trunc_requires_nanos(granularity: &str, has_timezone: bool) -> bool { + match granularity { + "microsecond" | "millisecond" | "second" | "minute" => false, + "hour" | "day" => has_timezone, + "week" | "month" | "quarter" | "year" => true, + _ => unreachable!("granularity was normalized before compatibility dispatch"), + } +} + +/// Returns whether a timezone has a zero UTC offset at every instant. +/// +/// Keep this list conservative: an unlisted timezone takes the zone-aware path, which is always +/// correct even when its current offset happens to be zero. +fn is_utc_timezone(timezone: &str) -> bool { + matches!( + timezone, + "UTC" | "Etc/UTC" | "Etc/GMT" | "GMT" | "Z" | "+00:00" | "-00:00" | "00:00" + ) +} + +/// Truncate timezone-aware timestamps to the local minute boundary. +/// +/// DataFusion floors the stored UTC microseconds directly. That is only equivalent to Spark's +/// local-time truncation when the timezone offset is a whole number of minutes. Historical offsets +/// can include seconds, so resolve the offset for each instant before finding the local remainder. +/// If truncation crosses an offset transition, re-resolve the truncated local datetime just like +/// `ZonedDateTime.truncatedTo` rather than retaining the input instant's offset. +fn timestamp_trunc_minute_tz( + array: &TimestampMicrosecondArray, + timezone: &str, +) -> Result { + as_timestamp_tz_with_op::<&TimestampMicrosecondArray, TimestampMicrosecondType, _>( + ArrayIter::new(array), + TimestampMicrosecondBuilder::with_capacity(array.len()), + timezone, + |dt| { + let micros = dt.timestamp_micros(); + let timezone = dt.timezone(); + let original_offset_secs = dt.offset().fix().local_minus_utc(); + let offset_micros = i64::from(original_offset_secs) * 1_000_000; + let candidate = micros - (micros + offset_micros).rem_euclid(MICROS_PER_MINUTE); + let candidate_dt = + as_datetime_with_timezone::(candidate, timezone) + .expect("truncated minute candidate must be a valid datetime"); + let candidate_offset_secs = candidate_dt.offset().fix().local_minus_utc(); + + if candidate_offset_secs == original_offset_secs { + return candidate; + } + + let truncated_local = dt + .naive_local() + .with_second(0) + .and_then(|local| local.with_nanosecond(0)) + .expect("truncated local minute must be a valid datetime"); + match timezone.from_local_datetime(&truncated_local) { + LocalResult::Single(resolved) => resolved.timestamp_micros(), + LocalResult::Ambiguous(earlier, later) => { + // ZonedDateTime retains the original offset when it is valid in an overlap. + if earlier.offset().fix().local_minus_utc() == original_offset_secs { + earlier.timestamp_micros() + } else if later.offset().fix().local_minus_utc() == original_offset_secs { + later.timestamp_micros() + } else { + earlier.timestamp_micros() + } + } + LocalResult::None => { + // The candidate lies immediately before a forward transition. Java advances + // a nonexistent local time by the gap, which is equivalent to resolving it + // with the candidate's pre-transition offset. + naive_to_micros(truncated_local) - i64::from(candidate_offset_secs) * 1_000_000 + } + } + }, + ) +} + +/// Truncate timezone-aware timestamps to a coarse local-date boundary. +/// +/// Spark resolves these boundaries with `LocalDate.atStartOfDay`. Unlike +/// `ZonedDateTime.truncatedTo`, an overlap always selects the earlier instant, independent of +/// the input timestamp's offset. A midnight in a gap resolves to the transition instant, +/// the first valid local time after the gap, even when the gap starts before midnight. +fn timestamp_trunc_coarse_tz( + array: &TimestampMicrosecondArray, + format: &str, + timezone: &str, +) -> Result { + let trunc_fn = ntz_trunc_fn_for_format(format)?; + let tz: Tz = timezone.parse()?; + + as_timestamp_tz_with_op::<&TimestampMicrosecondArray, TimestampMicrosecondType, _>( + ArrayIter::new(array), + TimestampMicrosecondBuilder::with_capacity(array.len()), + timezone, + |dt| { + let truncated_local = trunc_fn(dt.naive_local()) + .expect("truncated coarse local datetime must be a valid datetime"); + + match tz.from_local_datetime(&truncated_local) { + LocalResult::Single(resolved) => resolved.timestamp_micros(), + // `LocalDate.atStartOfDay` takes the earlier occurrence in an overlap. + LocalResult::Ambiguous(earlier, _) => earlier.timestamp_micros(), + LocalResult::None => { + // atStartOfDay returns the end of the gap, not midnight shifted by its + // length. Find a valid pre-gap instant, extending the usual three-hour + // probe for zones that skipped an entire date (e.g. Pacific/Apia). + let mut probe = truncated_local - Duration::hours(3); + let before_gap = loop { + match tz.from_local_datetime(&probe) { + LocalResult::Single(resolved) | LocalResult::Ambiguous(resolved, _) => { + break resolved; + } + LocalResult::None => probe -= Duration::hours(3), + } + }; + let pre_gap_offset_secs = before_gap.offset().fix().local_minus_utc(); + let mut low = before_gap.timestamp(); + let mut high = + truncated_local.and_utc().timestamp() - i64::from(pre_gap_offset_secs); + // TZ offsets and transitions have whole-second precision. The old offset + // holds at low, and has changed by high; find the first changed second. + while high - low > 1 { + let mid = low + (high - low) / 2; + let utc = DateTime::from_timestamp(mid, 0) + .expect("midnight gap probe must be a valid datetime") + .naive_utc(); + if tz.offset_from_utc_datetime(&utc).fix().local_minus_utc() + == pre_gap_offset_secs + { + low = mid; + } else { + high = mid; + } + } + high * 1_000_000 + } } }, - None => builder.append_null(), + ) +} + +/// Handle the rare fine-unit batch whose physical values would underflow DataFusion's +/// unchecked subtraction. Iterate logical values so NULL slots never cause errors. +fn timestamp_trunc_fine_boundary( + array: &TimestampMicrosecondArray, + granularity: &str, + unit: i64, + wrap_second_millisecond_overflow: bool, +) -> Result { + let mut builder = TimestampMicrosecondBuilder::with_capacity(array.len()); + for value in array.iter() { + match value { + None => builder.append_null(), + Some(micros) => { + let remainder = micros.rem_euclid(unit); + let truncated = if wrap_second_millisecond_overflow + && matches!(granularity, "second" | "millisecond") + { + // Before 4.2, Spark uses unchecked Long subtraction for these + // timezone-independent units. + micros.wrapping_sub(remainder) + } else { + // MINUTE/HOUR/DAY use Spark's exact instant-to-microseconds conversion, and + // Spark 4.2 and later (SPARK-56663) check the SECOND/MILLISECOND + // subtraction too. + micros.checked_sub(remainder).ok_or_else(|| { + SparkError::Internal(format!( + "long overflow: Timestamp {micros} out of range after date_trunc({granularity})" + )) + })? + }; + builder.append_value(truncated); + } + } + } + Ok(builder.finish().with_timezone_opt(array.timezone())) +} + +fn timestamp_trunc_upstream( + array: &TimestampMicrosecondArray, + format: &str, + wrap_second_millisecond_overflow: bool, +) -> Result { + let granularity = normalize_timestamp_trunc_format(format)?; + + if matches!( + granularity, + "hour" | "day" | "week" | "month" | "quarter" | "year" + ) && array.timezone().is_some_and(is_utc_timezone) + { + // UTC local time equals the instant. Removing its label lets DataFusion use arithmetic + // for HOUR/DAY (without a nanosecond limit) and integer calendar code for coarse units. + // The unlabelled call still applies the coarse-unit range guard and calendar fallback. + let input = array.clone().with_timezone_opt(None::>); + let result = timestamp_trunc_upstream(&input, format, wrap_second_millisecond_overflow)?; + return Ok(result.with_timezone_opt(array.timezone())); + } + + if granularity == "minute" { + if let Some(timezone) = array.timezone().filter(|tz| !is_utc_timezone(tz)) { + return timestamp_trunc_minute_tz(array, timezone); + } + } + + if matches!(granularity, "week" | "month" | "quarter" | "year") { + if let Some(timezone) = array.timezone().filter(|tz| !is_utc_timezone(tz)) { + return timestamp_trunc_coarse_tz(array, format, timezone); + } + } + + let fine_unit = match granularity { + "millisecond" => Some(1_000), + "second" => Some(1_000_000), + "minute" => Some(MICROS_PER_MINUTE), + "hour" if array.timezone().is_none() => Some(3_600_000_000), + "day" if array.timezone().is_none() => Some(MICROS_PER_DAY), + _ => None, + }; + if let Some(unit) = fine_unit { + // First aligned microsecond value that can be represented after truncation. + let lower = i64::MIN + (unit - i64::MIN.rem_euclid(unit)).rem_euclid(unit); + // Check physical values too: DataFusion processes NULL slots before restoring validity. + if array.values().iter().any(|micros| *micros < lower) { + return timestamp_trunc_fine_boundary( + array, + granularity, + unit, + wrap_second_millisecond_overflow, + ); + } + } + + let requires_nanos = + datafusion_timestamp_trunc_requires_nanos(granularity, array.timezone().is_some()); + + if !requires_nanos + || array + .iter() + .flatten() + .all(fits_datafusion_coarse_trunc_range) + { + let result = datafusion_date_trunc(Arc::new(array.clone()), granularity)?; + return Ok(result + .as_any() + .downcast_ref::() + .expect("DataFusion date_trunc timestamp result mismatch") + .clone()); + } + + let upstream_input = TimestampMicrosecondArray::from_iter( + array + .iter() + .map(|value| value.filter(|micros| fits_datafusion_coarse_trunc_range(*micros))), + ) + .with_timezone_opt(array.timezone()); + let upstream = datafusion_date_trunc(Arc::new(upstream_input), granularity)?; + let upstream = upstream + .as_any() + .downcast_ref::() + .expect("DataFusion date_trunc timestamp result mismatch"); + // Non-UTC HOUR/DAY still need the zone-aware fallback. UTC coarse units have already + // removed their timezone label and must not fall back through chrono. + let legacy = if array.timezone().is_some() { + let input = TimestampMicrosecondArray::from_iter( + array + .iter() + .map(|value| value.filter(|micros| !fits_datafusion_coarse_trunc_range(*micros))), + ) + .with_timezone_opt(array.timezone()); + Some(timestamp_trunc_legacy(&input, format)?) + } else { + None + }; + let mut builder = TimestampMicrosecondBuilder::with_capacity(array.len()); + for (index, value) in array.iter().enumerate() { + match value { + None => builder.append_null(), + Some(micros) if fits_datafusion_coarse_trunc_range(micros) => { + builder.append_option((!upstream.is_null(index)).then(|| upstream.value(index))); + } + Some(micros) => match &legacy { + Some(result) => { + builder.append_option((!result.is_null(index)).then(|| result.value(index))); + } + None => builder.append_value(timestamp_trunc_coarse_micros(micros, granularity)?), + }, + } } - Ok(()) + Ok(builder.finish().with_timezone_opt(array.timezone())) } pub(crate) fn timestamp_trunc( array: &PrimitiveArray, format: String, + wrap_second_millisecond_overflow: bool, ) -> Result where T: ArrowTemporalType + ArrowNumericType, i64: From, { - let builder = TimestampMicrosecondBuilder::with_capacity(array.len()); - let iter = ArrayIter::new(array); match array.data_type() { - DataType::Timestamp(TimeUnit::Microsecond, None) => { - // TimestampNTZ: operate directly on naive microsecond values without timezone - timestamp_trunc_ntz(array, format) - } - DataType::Timestamp(TimeUnit::Microsecond, Some(tz)) => { - let trunc_fn = tz_trunc_fn_for_format(&format)?; - as_timestamp_tz_with_op::<&PrimitiveArray, T, _>(iter, builder, tz, |dt| { - as_micros_from_unix_epoch_utc(trunc_fn(dt)) - }) - } + DataType::Timestamp(TimeUnit::Microsecond, _) => timestamp_trunc_upstream( + array + .as_any() + .downcast_ref::() + .expect("TimestampMicrosecond type mismatch"), + &format, + wrap_second_millisecond_overflow, + ), dt => return_compute_error_with!( "Unsupported input type '{:?}' for function 'timestamp_trunc'", dt @@ -740,194 +1253,372 @@ where /// /// format is an array of strings specifying the format to apply to the corresponding timestamp /// value. The array may be a dictionary array. +/// +/// wrap_second_millisecond_overflow is passed to the literal-format kernel, as in +/// `timestamp_trunc_dyn`. pub(crate) fn timestamp_trunc_array_fmt_dyn( array: &dyn Array, formats: &dyn Array, + wrap_second_millisecond_overflow: bool, ) -> Result { - match (array.data_type().clone(), formats.data_type().clone()) { - (DataType::Dictionary(_, _), DataType::Dictionary(_, _)) => { - downcast_dictionary_array!( - formats => { - downcast_dictionary_array!( - array => { - timestamp_trunc_array_fmt_dict_dict( - &array.downcast_dict::().unwrap(), - &formats.downcast_dict::().unwrap()) - .map(|a| Arc::new(a) as ArrayRef) - } - dt => return_compute_error_with!("timestamp_trunc does not support", dt) - ) - } - fmt => return_compute_error_with!("timestamp_trunc does not support format type", fmt), - ) + let array = unpack_dictionary(array)?; + let formats = unpack_dictionary(formats)?; + let array = match array.data_type() { + DataType::Timestamp(TimeUnit::Microsecond, _) => { + array.as_primitive::() } - (DataType::Dictionary(_, _), DataType::Utf8) => { - downcast_dictionary_array!( - array => { - timestamp_trunc_array_fmt_dict_plain( - &array.downcast_dict::>().unwrap(), - formats.as_any().downcast_ref::() - .expect("Unexpected value type in formats")) - .map(|a| Arc::new(a) as ArrayRef) - } - dt => return_compute_error_with!("timestamp_trunc does not support", dt), - ) + dt => { + return_compute_error_with!("Unsupported input type for function 'timestamp_trunc'", dt) } - (DataType::Timestamp(TimeUnit::Microsecond, _), DataType::Dictionary(_, _)) => { - downcast_dictionary_array!( - formats => { - downcast_temporal_array!(array => { - timestamp_trunc_array_fmt_plain_dict( - array, - &formats.downcast_dict::().unwrap()) - .map(|a| Arc::new(a) as ArrayRef) - } - dt => return_compute_error_with!("timestamp_trunc does not support", dt), - ) - } - fmt => return_compute_error_with!("timestamp_trunc does not support format type", fmt), - ) + }; + let formats = match formats.data_type() { + DataType::Utf8 => formats.as_string::(), + fmt => return_compute_error_with!("timestamp_trunc does not support format type", fmt), + }; + Ok(Arc::new(timestamp_trunc_by_row_format( + array, + formats, + wrap_second_millisecond_overflow, + )?)) +} + +fn unpack_dictionary(array: &dyn Array) -> Result { + match array.data_type() { + DataType::Dictionary(_, value_type) => arrow::compute::cast(array, value_type) + .map_err(|error| SparkError::Internal(error.to_string())), + _ => Ok(make_array(array.to_data())), + } +} + +/// Truncates each row with the format in the same row. +/// +/// Rows are grouped by format, and each group goes through the literal-format kernel with the +/// other rows set to NULL. A row is therefore truncated by exactly the rules a literal format +/// applies, including the DST handling, and a value in one group cannot make another group fail. +/// A NULL format gives a NULL result, as in Spark. +fn timestamp_trunc_by_row_format( + array: &TimestampMicrosecondArray, + formats: &StringArray, + wrap_second_millisecond_overflow: bool, +) -> Result { + if array.len() != formats.len() { + return Err(SparkError::Internal(format!( + "timestamp_trunc has {} values but {} formats", + array.len(), + formats.len() + ))); + } + + let mut granularities: Vec<&'static str> = Vec::new(); + // A batch rarely holds more than a few spellings, so parse each one once. + let mut spellings: Vec<(&str, usize)> = Vec::new(); + let mut row_groups: Vec> = Vec::with_capacity(array.len()); + let mut null_format_hides_value = false; + for index in 0..array.len() { + if array.is_null(index) { + row_groups.push(None); + continue; } - (DataType::Timestamp(TimeUnit::Microsecond, _), DataType::Utf8) => { - downcast_temporal_array!( - array => { - timestamp_trunc_array_fmt_plain_plain(array, - formats.as_any().downcast_ref::().expect("Unexpected value type in formats")) - .map(|a| Arc::new(a) as ArrayRef) - }, - dt => return_compute_error_with!("timestamp_trunc does not support", dt), + if formats.is_null(index) { + null_format_hides_value = true; + row_groups.push(None); + continue; + } + let spelling = formats.value(index); + let group = match spellings.iter().find(|(seen, _)| *seen == spelling) { + Some((_, group)) => *group, + None => { + let granularity = normalize_timestamp_trunc_format(spelling)?; + let group = match granularities.iter().position(|g| *g == granularity) { + Some(group) => group, + None => { + granularities.push(granularity); + granularities.len() - 1 + } + }; + spellings.push((spelling, group)); + group + } + }; + row_groups.push(Some(group)); + } + + if granularities.len() == 1 && !null_format_hides_value { + return timestamp_trunc_upstream(array, granularities[0], wrap_second_millisecond_overflow); + } + + let truncated = granularities + .iter() + .enumerate() + .map(|(group, granularity)| { + let outside_group: BooleanArray = row_groups + .iter() + .map(|row_group| Some(*row_group != Some(group))) + .collect(); + let input = arrow::compute::nullif(array, &outside_group) + .map_err(|error| SparkError::Internal(error.to_string()))?; + timestamp_trunc_upstream( + input.as_primitive(), + granularity, + wrap_second_millisecond_overflow, ) + }) + .collect::, _>>()?; + + let mut builder = TimestampMicrosecondBuilder::with_capacity(array.len()); + for (index, row_group) in row_groups.iter().enumerate() { + match row_group.map(|group| &truncated[group]) { + Some(result) if result.is_valid(index) => builder.append_value(result.value(index)), + _ => builder.append_null(), } - (dt, fmt) => Err(SparkError::Internal(format!( - "Unsupported datatype: {dt:}, format: {fmt:?} for function 'timestamp_trunc'" - ))), } + Ok(builder.finish().with_timezone_opt(array.timezone())) } -macro_rules! timestamp_trunc_array_fmt_helper { - ($array: ident, $formats: ident, $datatype: ident) => {{ - let mut builder = TimestampMicrosecondBuilder::with_capacity($array.len()); - let iter = $array.into_iter(); +#[cfg(test)] +mod tests { + use super::{ + naive_to_micros, normalize_timestamp_trunc_format, ntz_trunc_fn_for_format, + timestamp_trunc_coarse_micros, timestamp_trunc_dictionary_needs_mask, timestamp_trunc_ntz, + MICROS_PER_DAY, TIMESTAMP_TRUNC_ALIASES, + }; + use crate::kernels::temporal::{ + date_trunc, date_trunc_array_fmt_dyn, date_trunc_dyn, timestamp_trunc, + timestamp_trunc_array_fmt_dyn, timestamp_trunc_dyn, + }; + use crate::SparkError; + use arrow::array::{ + builder::{PrimitiveDictionaryBuilder, StringDictionaryBuilder}, + iterator::ArrayIter, + types::{Date32Type, Int32Type, TimestampMicrosecondType}, + Array, Date32Array, DictionaryArray, Int32Array, PrimitiveArray, StringArray, + TimestampMicrosecondArray, + }; + use arrow::datatypes::{DataType, TimeUnit}; + use chrono::{DateTime, Datelike, NaiveDate}; + use std::sync::Arc; + + fn epoch_days(date: &str) -> i32 { + NaiveDate::parse_from_str(date, "%Y-%m-%d") + .unwrap() + .num_days_from_ce() + - 719_163 + } + + fn assert_date_trunc(format: &str, input: &[Option<&str>], expected: &[Option<&str>]) { + let input = Date32Array::from( + input + .iter() + .map(|date| date.map(epoch_days)) + .collect::>(), + ); + let expected = Date32Array::from( + expected + .iter() + .map(|date| date.map(epoch_days)) + .collect::>(), + ); + assert_eq!(date_trunc(&input, format.to_string()).unwrap(), expected); + } + + fn instant_micros(instant: &str) -> i64 { + DateTime::parse_from_rfc3339(instant) + .unwrap() + .timestamp_micros() + } + + fn assert_timestamp_trunc( + format: &str, + timezone: Option<&str>, + input: &[Option<&str>], + expected: &[Option<&str>], + ) { + let input = TimestampMicrosecondArray::from( + input + .iter() + .map(|instant| instant.map(instant_micros)) + .collect::>(), + ) + .with_timezone_opt(timezone); + let expected = TimestampMicrosecondArray::from( + expected + .iter() + .map(|instant| instant.map(instant_micros)) + .collect::>(), + ) + .with_timezone_opt(timezone); assert_eq!( - $array.len(), - $formats.len(), - "lengths of values array and format array must be the same" + timestamp_trunc(&input, format.to_string(), true).unwrap(), + expected ); - match $datatype { - DataType::Timestamp(TimeUnit::Microsecond, None) => { - // TimestampNTZ: operate directly on naive microsecond values - for (index, val) in iter.enumerate() { - let micros_val = val.map(|v| i64::from(v)); - let trunc_fn = ntz_trunc_fn_for_format($formats.value(index))?; - timestamp_trunc_ntz_single(micros_val, &mut builder, trunc_fn)?; - } - Ok(builder.finish()) - } - DataType::Timestamp(TimeUnit::Microsecond, Some(tz_str)) => { - let tz: Tz = tz_str.parse()?; - for (index, val) in iter.enumerate() { - let trunc_fn = tz_trunc_fn_for_format($formats.value(index))?; - as_timestamp_tz_with_op_single::(val, &mut builder, &tz, |dt| { - as_micros_from_unix_epoch_utc(trunc_fn(dt)) - })?; - } - Ok(builder.finish().with_timezone(tz_str.as_ref())) - } - dt => { - return_compute_error_with!( - "Unsupported input type '{:?}' for function 'timestamp_trunc'", - dt - ) - } + } + + #[test] + fn test_date_trunc() { + for format in ["YEAR", "YYYY", "YY", "year", "Year", "yEaR"] { + assert_date_trunc(format, &[Some("2024-05-17")], &[Some("2024-01-01")]); + } + for format in ["MONTH", "MON", "MM", "month", "Mon"] { + assert_date_trunc(format, &[Some("2024-05-17")], &[Some("2024-05-01")]); } - }}; -} -fn timestamp_trunc_array_fmt_plain_plain( - array: &PrimitiveArray, - formats: &StringArray, -) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: From, -{ - let data_type = array.data_type(); - timestamp_trunc_array_fmt_helper!(array, formats, data_type) -} -fn timestamp_trunc_array_fmt_plain_dict( - array: &PrimitiveArray, - formats: &TypedDictionaryArray, -) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: From, - K: ArrowDictionaryKeyType, -{ - let data_type = array.data_type(); - timestamp_trunc_array_fmt_helper!(array, formats, data_type) -} + assert_date_trunc( + "QUARTER", + &[ + Some("2024-01-01"), + Some("2024-03-31"), + Some("2024-04-01"), + Some("2024-06-30"), + Some("2024-07-01"), + Some("2024-10-01"), + ], + &[ + Some("2024-01-01"), + Some("2024-01-01"), + Some("2024-04-01"), + Some("2024-04-01"), + Some("2024-07-01"), + Some("2024-10-01"), + ], + ); -fn timestamp_trunc_array_fmt_dict_plain( - array: &TypedDictionaryArray>, - formats: &StringArray, -) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: From, - K: ArrowDictionaryKeyType, -{ - let data_type = array.values().data_type(); - timestamp_trunc_array_fmt_helper!(array, formats, data_type) -} + assert_date_trunc( + "week", + &[ + Some("2024-05-13"), + Some("2024-05-14"), + Some("2024-05-19"), + Some("2024-05-01"), + Some("2024-01-01"), + Some("2023-12-31"), + ], + &[ + Some("2024-05-13"), + Some("2024-05-13"), + Some("2024-05-13"), + Some("2024-04-29"), + Some("2024-01-01"), + Some("2023-12-25"), + ], + ); -fn timestamp_trunc_array_fmt_dict_dict( - array: &TypedDictionaryArray>, - formats: &TypedDictionaryArray, -) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: From, - K: ArrowDictionaryKeyType, - F: ArrowDictionaryKeyType, -{ - let data_type = array.values().data_type(); - timestamp_trunc_array_fmt_helper!(array, formats, data_type) -} + let dates = [ + Some("2024-02-29"), + Some("2000-02-29"), + Some("1900-02-28"), + Some("1969-12-31"), + Some("1960-02-29"), + Some("1900-01-01"), + // The input fits TimestampNanosecond, but truncating it to YEAR does not. + Some("1677-09-22"), + // Valid Spark Date32 outside TimestampNanosecond's range. DataFusion 55.1 + // date_trunc internally converts coarse granularities to nanoseconds. + Some("3333-05-17"), + None, + ]; + assert_date_trunc( + "YEAR", + &dates, + &[ + Some("2024-01-01"), + Some("2000-01-01"), + Some("1900-01-01"), + Some("1969-01-01"), + Some("1960-01-01"), + Some("1900-01-01"), + Some("1677-01-01"), + Some("3333-01-01"), + None, + ], + ); + assert_date_trunc( + "QUARTER", + &dates, + &[ + Some("2024-01-01"), + Some("2000-01-01"), + Some("1900-01-01"), + Some("1969-10-01"), + Some("1960-01-01"), + Some("1900-01-01"), + Some("1677-07-01"), + Some("3333-04-01"), + None, + ], + ); + assert_date_trunc( + "MONTH", + &dates, + &[ + Some("2024-02-01"), + Some("2000-02-01"), + Some("1900-02-01"), + Some("1969-12-01"), + Some("1960-02-01"), + Some("1900-01-01"), + Some("1677-09-01"), + Some("3333-05-01"), + None, + ], + ); + assert_date_trunc( + "WEEK", + &dates, + &[ + Some("2024-02-26"), + Some("2000-02-28"), + Some("1900-02-26"), + Some("1969-12-29"), + Some("1960-02-29"), + Some("1900-01-01"), + Some("1677-09-20"), + Some("3333-05-11"), + None, + ], + ); -#[cfg(test)] -mod tests { - use crate::kernels::temporal::{ - date_trunc, date_trunc_array_fmt_dyn, timestamp_trunc, timestamp_trunc_array_fmt_dyn, - }; - use arrow::array::{ - builder::{PrimitiveDictionaryBuilder, StringDictionaryBuilder}, - iterator::ArrayIter, - types::{Date32Type, Int32Type, TimestampMicrosecondType}, - Array, Date32Array, PrimitiveArray, StringArray, TimestampMicrosecondArray, - }; - use std::sync::Arc; + let input = Date32Array::from(vec![epoch_days("2024-05-17")]); + for format in ["DAY", "HOUR", "SECOND", "invalid", " YEAR ", ""] { + let SparkError::Internal(message) = date_trunc(&input, format.to_string()).unwrap_err() + else { + panic!("expected an internal unsupported-format error"); + }; + assert_eq!( + message, + format!("Unsupported format: {format:?} for function 'date_trunc'") + ); + } + } #[test] - #[cfg_attr(miri, ignore)] // test takes too long with miri - fn test_date_trunc() { - let size = 1000; - let mut vec: Vec = Vec::with_capacity(size); - for i in 0..size { - vec.push(i as i32); - } - let array = Date32Array::from(vec); - for fmt in [ - "YEAR", "YYYY", "YY", "QUARTER", "MONTH", "MON", "MM", "WEEK", - ] { - match date_trunc(&array, fmt.to_string()) { - Ok(a) => { - for i in 0..size { - assert!(array.values().get(i) >= a.values().get(i)) - } - } - _ => unreachable!(), - } - } + fn test_date_trunc_scalar_format_dictionary() { + let mut builder = PrimitiveDictionaryBuilder::::new(); + builder.append(epoch_days("2024-05-17")).unwrap(); + builder.append(epoch_days("2024-06-30")).unwrap(); + builder.append(epoch_days("2024-05-17")).unwrap(); + builder.append_null(); + builder.append(epoch_days("1969-12-31")).unwrap(); + let input = builder.finish(); + let input_keys = input.keys().clone(); + + let result = date_trunc_dyn(&input, "MONTH".to_string()).unwrap(); + let result = result + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(result.keys(), &input_keys); + + let decoded = result.downcast_dict::().unwrap(); + assert_eq!( + decoded.into_iter().collect::>(), + vec![ + Some(epoch_days("2024-05-01")), + Some(epoch_days("2024-06-01")), + Some(epoch_days("2024-05-01")), + None, + Some(epoch_days("1969-12-01")), + ] + ); } #[test] @@ -1047,42 +1738,1011 @@ mod tests { } #[test] - #[cfg_attr(miri, ignore)] // test takes too long with miri fn test_timestamp_trunc() { - let size = 1000; - let mut vec: Vec = Vec::with_capacity(size); - for i in 0..size { - vec.push(i as i64); + let input = [Some("2024-05-17T12:34:56.123456Z"), None]; + for format in ["YEAR", "YYYY", "YY", "year", "Year", "yEaR"] { + assert_timestamp_trunc( + format, + Some("UTC"), + &input, + &[Some("2024-01-01T00:00:00Z"), None], + ); } - let array = TimestampMicrosecondArray::from(vec).with_timezone_utc(); - for fmt in [ + for format in ["MONTH", "MON", "MM", "month", "Mon"] { + assert_timestamp_trunc( + format, + Some("UTC"), + &input, + &[Some("2024-05-01T00:00:00Z"), None], + ); + } + for (format, expected) in [ + ("QUARTER", "2024-04-01T00:00:00Z"), + ("WEEK", "2024-05-13T00:00:00Z"), + ("DAY", "2024-05-17T00:00:00Z"), + ("DD", "2024-05-17T00:00:00Z"), + ("HOUR", "2024-05-17T12:00:00Z"), + ("MINUTE", "2024-05-17T12:34:00Z"), + ("SECOND", "2024-05-17T12:34:56Z"), + ("MILLISECOND", "2024-05-17T12:34:56.123Z"), + ("MICROSECOND", "2024-05-17T12:34:56.123456Z"), + ] { + assert_timestamp_trunc(format, Some("UTC"), &input, &[Some(expected), None]); + } + + let invalid_input = + TimestampMicrosecondArray::from(vec![instant_micros("2024-05-17T12:34:56Z")]) + .with_timezone_utc(); + for format in ["MILLISECONDS", "invalid", " DAY ", ""] { + let SparkError::Internal(message) = + timestamp_trunc(&invalid_input, format.to_string(), true).unwrap_err() + else { + panic!("expected an internal unsupported-format error"); + }; + assert_eq!( + message, + format!("Unsupported format: {format:?} for function 'timestamp_trunc'") + ); + } + } + + #[test] + fn test_timestamp_trunc_utc_aliases_and_wide_range() { + let input = [ + Some("1500-06-15T12:34:56.123456Z"), + Some("1678-06-01T12:34:56.123456Z"), + Some("1969-12-31T23:59:59.123456Z"), + Some("3333-05-17T12:34:56.123456Z"), + None, + ]; + for timezone in [ + "UTC", "Etc/UTC", "Etc/GMT", "GMT", "Z", "+00:00", "-00:00", "00:00", + ] { + for (format, expected) in [ + ( + "HOUR", + [ + Some("1500-06-15T12:00:00Z"), + Some("1678-06-01T12:00:00Z"), + Some("1969-12-31T23:00:00Z"), + Some("3333-05-17T12:00:00Z"), + None, + ], + ), + ( + "DAY", + [ + Some("1500-06-15T00:00:00Z"), + Some("1678-06-01T00:00:00Z"), + Some("1969-12-31T00:00:00Z"), + Some("3333-05-17T00:00:00Z"), + None, + ], + ), + ( + "YEAR", + [ + Some("1500-01-01T00:00:00Z"), + Some("1678-01-01T00:00:00Z"), + Some("1969-01-01T00:00:00Z"), + Some("3333-01-01T00:00:00Z"), + None, + ], + ), + ] { + assert_timestamp_trunc(format, Some(timezone), &input, &expected); + } + } + } + + #[test] + fn test_timestamp_trunc_wide_range_fallback() { + let input = [ + Some("2024-05-17T12:34:56.123456Z"), + Some("3333-05-17T12:34:56.123456Z"), + Some("1969-12-31T23:59:59.123456Z"), + Some("1677-09-22T00:00:00Z"), + None, + ]; + assert_timestamp_trunc( "YEAR", - "YYYY", - "YY", + Some("UTC"), + &input, + &[ + Some("2024-01-01T00:00:00Z"), + Some("3333-01-01T00:00:00Z"), + Some("1969-01-01T00:00:00Z"), + Some("1677-01-01T00:00:00Z"), + None, + ], + ); + assert_timestamp_trunc( "QUARTER", - "MONTH", - "MON", - "MM", + Some("UTC"), + &input, + &[ + Some("2024-04-01T00:00:00Z"), + Some("3333-04-01T00:00:00Z"), + Some("1969-10-01T00:00:00Z"), + Some("1677-07-01T00:00:00Z"), + None, + ], + ); + assert_timestamp_trunc( "WEEK", - "DAY", - "DD", + Some("UTC"), + &input, + &[ + Some("2024-05-13T00:00:00Z"), + Some("3333-05-11T00:00:00Z"), + Some("1969-12-29T00:00:00Z"), + Some("1677-09-20T00:00:00Z"), + None, + ], + ); + } + + #[test] + fn test_timestamp_trunc_out_of_chrono_range() { + for timezone in [None, Some("UTC"), Some("Etc/UTC"), Some("+00:00")] { + let input = TimestampMicrosecondArray::from(vec![ + Some(instant_micros("2024-05-17T12:34:56Z")), + Some(i64::MAX), + Some(i64::MAX - 1), + Some(i64::MAX - 2), + Some(instant_micros("3333-05-17T12:34:56Z")), + Some(-9_000_000_000_000_000_000), + None, + ]) + .with_timezone_opt(timezone); + // Expected extremes were computed independently with java.time.LocalDate: + // MAX, MAX - 1, and MAX - 2 are all +294247-01-10; + // the negative input is -283229-05-10. + for (format, recent, future, maximum, negative) in [ + ( + "YEAR", + "2024-01-01T00:00:00Z", + "3333-01-01T00:00:00Z", + 9_223_371_244_800_000_000, + -9_000_011_174_400_000_000, + ), + ( + "QUARTER", + "2024-04-01T00:00:00Z", + "3333-04-01T00:00:00Z", + 9_223_371_244_800_000_000, + -9_000_003_398_400_000_000, + ), + ( + "MONTH", + "2024-05-01T00:00:00Z", + "3333-05-01T00:00:00Z", + 9_223_371_244_800_000_000, + -9_000_000_806_400_000_000, + ), + ( + "WEEK", + "2024-05-13T00:00:00Z", + "3333-05-11T00:00:00Z", + 9_223_371_504_000_000_000, + -9_000_000_028_800_000_000, + ), + ] { + let expected = TimestampMicrosecondArray::from(vec![ + Some(instant_micros(recent)), + Some(maximum), + Some(maximum), + Some(maximum), + Some(instant_micros(future)), + Some(negative), + None, + ]) + .with_timezone_opt(timezone); + assert_eq!( + timestamp_trunc(&input, format.to_string(), true).unwrap(), + expected + ); + for (alias, canonical) in TIMESTAMP_TRUNC_ALIASES { + if canonical == normalize_timestamp_trunc_format(format).unwrap() { + assert_eq!( + timestamp_trunc(&input, alias.to_lowercase(), true).unwrap(), + expected + ); + } + } + } + } + } + + #[test] + fn test_timestamp_trunc_coarse_overflow() { + for timezone in [None, Some("UTC"), Some("Etc/UTC"), Some("+00:00")] { + let input = + TimestampMicrosecondArray::from(vec![Some(i64::MIN)]).with_timezone_opt(timezone); + for format in ["YEAR", "QUARTER", "MONTH", "WEEK"] { + let error = timestamp_trunc(&input, format.to_string(), true).unwrap_err(); + assert!(error.to_string().contains("out of range"), "{error}"); + } + } + } + + #[test] + fn test_timestamp_trunc_integer_calendar_matches_chrono() { + // Exercise negative eras, year zero, Gregorian leap centuries, and month boundaries. + for year in [-2000, -400, -1, 0, 1, 1500, 1600, 1900, 2000, 2024, 3333] { + for month in 1..=12 { + for day in [1, 15, 28, 29, 30, 31] { + let Some(date) = NaiveDate::from_ymd_opt(year, month, day) else { + continue; + }; + let datetime = date.and_hms_micro_opt(23, 59, 59, 999_999).unwrap(); + for format in ["YEAR", "QUARTER", "MONTH", "WEEK"] { + let expected = naive_to_micros( + ntz_trunc_fn_for_format(format).unwrap()(datetime).unwrap(), + ); + assert_eq!( + timestamp_trunc_coarse_micros( + naive_to_micros(datetime), + normalize_timestamp_trunc_format(format).unwrap() + ) + .unwrap(), + expected + ); + } + } + } + } + } + + #[test] + fn test_timestamp_trunc_coarse_guard_boundaries() { + let lower = i64::MIN / 1_000 + 370 * MICROS_PER_DAY; + let upper = i64::MAX / 1_000; + let input = TimestampMicrosecondArray::from(vec![ + Some(lower - 1), + Some(lower), + Some(lower + 1), + Some(upper - 1), + Some(upper), + Some(upper + 1), + None, + ]); + for format in ["YEAR", "QUARTER", "MONTH", "WEEK"] { + let expected = timestamp_trunc_ntz(&input, format.to_string()).unwrap(); + assert_eq!( + timestamp_trunc(&input, format.to_string(), true).unwrap(), + expected + ); + } + } + + #[test] + fn test_timestamp_trunc_extreme_dictionary() { + let values = TimestampMicrosecondArray::from(vec![Some(i64::MAX), None]); + let keys = Int32Array::from(vec![Some(0), None, Some(1), Some(0)]); + let input = + DictionaryArray::::try_new(keys.clone(), Arc::new(values)) + .unwrap(); + let expected_values = + TimestampMicrosecondArray::from(vec![Some(9_223_371_244_800_000_000), None]); + let expected = DictionaryArray::::try_new( + keys, + Arc::new(expected_values), + ) + .unwrap(); + let result = timestamp_trunc_dyn(&input, "YEAR".to_string(), true).unwrap(); + assert_eq!(result.as_ref(), &expected as &dyn Array); + } + + #[test] + fn test_timestamp_trunc_empty_and_null_batches() { + for timezone in [None, Some("UTC")] { + for values in [vec![], vec![None, None]] { + let input = TimestampMicrosecondArray::from(values).with_timezone_opt(timezone); + for format in ["YEAR", "QUARTER", "MONTH", "WEEK"] { + assert_eq!( + timestamp_trunc(&input, format.to_string(), true).unwrap(), + input + ); + } + } + } + } + + #[test] + fn test_timestamp_trunc_dictionary_unused_overflow() { + for timezone in [None, Some("UTC")] { + let values = TimestampMicrosecondArray::from(vec![ + Some(instant_micros("2024-05-17T12:34:56Z")), + Some(i64::MIN), + None, + ]) + .with_timezone_opt(timezone); + for keys in [ + vec![Some(0), None, Some(2), Some(0)], + vec![None, None], + vec![], + vec![Some(0); 8192], + ] { + let keys = Int32Array::from(keys); + let input = + DictionaryArray::::try_new(keys.clone(), Arc::new(values.clone())) + .unwrap(); + let expected_values = TimestampMicrosecondArray::from(vec![ + Some(instant_micros("2024-01-01T00:00:00Z")), + None, + None, + ]) + .with_timezone_opt(timezone); + // Slice past the first key to check that referenced-entry masking respects offsets. + let input = input.slice( + usize::from(!input.is_empty()), + input.len().saturating_sub(1), + ); + let expected = + DictionaryArray::::try_new(keys, Arc::new(expected_values)).unwrap(); + let expected = expected.slice( + usize::from(!expected.is_empty()), + expected.len().saturating_sub(1), + ); + let result = timestamp_trunc_dyn(&input, "YEAR".into(), true).unwrap(); + // Unused entries can be masked, so compare the decoded logical rows. + let result = arrow::compute::cast( + result.as_ref(), + &DataType::Timestamp(TimeUnit::Microsecond, timezone.map(Into::into)), + ) + .unwrap(); + let expected = arrow::compute::cast(&expected, result.data_type()).unwrap(); + assert_eq!(result.as_ref(), expected.as_ref()); + let used = DictionaryArray::::try_new( + Int32Array::from(vec![1]), + Arc::new(values.clone()), + ) + .unwrap(); + assert!(timestamp_trunc_dyn(&used, "YEAR".into(), true) + .unwrap_err() + .to_string() + .contains("long overflow")); + } + } + } + + #[test] + fn test_timestamp_trunc_dictionary_unused_chrono_boundary() { + // These unused values can panic in timezone-aware calendar helpers, so they must still + // be masked before computing the distinct values, even in a low-cardinality dictionary. + for timezone in ["+01:00", "Asia/Tokyo"] { + let values = TimestampMicrosecondArray::from(vec![ + instant_micros("2024-05-17T12:34:56Z"), + chrono::NaiveDateTime::MIN.and_utc().timestamp_micros() + MICROS_PER_DAY, + ]) + .with_timezone(timezone); + let input = DictionaryArray::::try_new( + Int32Array::from(vec![Some(0); 8192]), + Arc::new(values.clone()), + ) + .unwrap(); + for format in ["YEAR", "WEEK"] { + let result = timestamp_trunc_dyn(&input, format.into(), true).unwrap(); + let decoded = arrow::compute::cast(result.as_ref(), values.data_type()).unwrap(); + let expected = timestamp_trunc(&values.slice(0, 1), format.into(), true).unwrap(); + let decoded = decoded + .as_any() + .downcast_ref::() + .unwrap(); + assert!(decoded.iter().all(|value| value == Some(expected.value(0)))); + } + } + } + + #[test] + fn test_timestamp_trunc_dictionary_fine_mask_decision() { + // SECOND/MILLISECOND ignore the timezone, so a dictionary of safe values is truncated + // through its values alone in either mode. Only a value that can fail needs the key scan. + for timezone in [None, Some("UTC"), Some("Asia/Tokyo")] { + let safe = TimestampMicrosecondArray::from(vec![ + Some(instant_micros("2024-05-17T12:34:56Z")), + None, + ]) + .with_timezone_opt(timezone); + let near_min = + TimestampMicrosecondArray::from(vec![Some(i64::MIN)]).with_timezone_opt(timezone); + for format in ["SECOND", "MILLISECOND"] { + for wrap in [true, false] { + assert!( + !timestamp_trunc_dictionary_needs_mask(&safe, 8192, format, wrap).unwrap() + ); + } + assert!( + !timestamp_trunc_dictionary_needs_mask(&near_min, 8192, format, true).unwrap() + ); + assert!( + timestamp_trunc_dictionary_needs_mask(&near_min, 8192, format, false).unwrap() + ); + } + } + } + + #[test] + fn test_timestamp_trunc_dictionary_fine_unused_extremes() { + // Unused MIN entries must not raise, whether SECOND/MILLISECOND wrap (Spark before 4.2) + // or fail (Spark 4.2 and later). + for wrap in [true, false] { + for timezone in [None, Some("UTC"), Some("Asia/Tokyo")] { + let values = TimestampMicrosecondArray::from(vec![ + Some(instant_micros("2024-05-17T12:34:56Z")), + Some(i64::MIN), + None, + ]) + .with_timezone_opt(timezone); + for keys in [ + vec![Some(0), None, Some(2), Some(0)], + vec![None, None], + vec![], + ] { + let input = DictionaryArray::::try_new( + Int32Array::from(keys), + Arc::new(values.clone()), + ) + .unwrap(); + let decoded_input = arrow::compute::cast(&input, values.data_type()).unwrap(); + for format in ["MICROSECOND", "MILLISECOND", "SECOND"] { + let result = timestamp_trunc_dyn(&input, format.into(), wrap).unwrap(); + let decoded = + arrow::compute::cast(result.as_ref(), values.data_type()).unwrap(); + let expected = + timestamp_trunc_dyn(decoded_input.as_ref(), format.into(), wrap) + .unwrap(); + assert_eq!(decoded.as_ref(), expected.as_ref()); + } + } + if !wrap { + // A used MIN entry still raises. + let used = DictionaryArray::::try_new( + Int32Array::from(vec![0, 1]), + Arc::new(values.clone()), + ) + .unwrap(); + for format in ["MILLISECOND", "SECOND"] { + assert!(timestamp_trunc_dyn(&used, format.into(), wrap) + .unwrap_err() + .to_string() + .contains("long overflow")); + } + } + } + } + } + + #[test] + fn test_timestamp_trunc_fine_poisoned_nulls() { + for wrap in [true, false] { + for timezone in [None, Some("UTC")] { + let input = TimestampMicrosecondArray::new( + vec![ + i64::MIN, + instant_micros("2024-05-17T12:34:56Z"), + i64::MIN + 1, + i64::MAX, + ] + .into(), + Some(arrow::buffer::NullBuffer::from(vec![ + false, true, false, true, + ])), + ) + .with_timezone_opt(timezone); + let clean = TimestampMicrosecondArray::from(vec![ + None, + Some(instant_micros("2024-05-17T12:34:56Z")), + None, + Some(i64::MAX), + ]) + .with_timezone_opt(timezone); + for format in [ + "MINUTE", + "HOUR", + "DAY", + "SECOND", + "MILLISECOND", + "MICROSECOND", + ] { + assert_eq!( + timestamp_trunc(&input, format.into(), wrap).unwrap(), + timestamp_trunc(&clean, format.into(), wrap).unwrap() + ); + let slice = input.slice(2, 2); + assert_eq!( + timestamp_trunc(&slice, format.into(), wrap).unwrap(), + timestamp_trunc(&clean.slice(2, 2), format.into(), wrap).unwrap() + ); + } + } + } + } + + #[test] + fn test_timestamp_trunc_fine_lower_bound() { + // First representable UTC boundary for each unit. + for (format, boundary) in [ + ("MINUTE", -9_223_372_036_800_000_000), + ("HOUR", -9_223_372_036_800_000_000), + ("DAY", -9_223_372_022_400_000_000), + ] { + for timezone in [None, Some("UTC")] { + for micros in [i64::MIN, i64::MIN + 1, boundary - 1] { + let input = + TimestampMicrosecondArray::from(vec![micros]).with_timezone_opt(timezone); + assert!(timestamp_trunc(&input, format.into(), true) + .unwrap_err() + .to_string() + .contains("long overflow")); + } + let input = TimestampMicrosecondArray::from(vec![boundary, boundary + 1]) + .with_timezone_opt(timezone); + let expected = TimestampMicrosecondArray::from(vec![boundary, boundary]) + .with_timezone_opt(timezone); + assert_eq!( + timestamp_trunc(&input, format.into(), true).unwrap(), + expected + ); + } + } + } + + #[test] + fn test_timestamp_trunc_fine_spark_wrapping() { + // Before 4.2, Spark uses unchecked Long subtraction for SECOND/MILLISECOND, even at MIN. + for (format, expected) in [ + ("SECOND", 9_223_372_036_854_551_616), + ("MILLISECOND", 9_223_372_036_854_775_616), + ] { + for timezone in [None, Some("UTC"), Some("Asia/Kolkata")] { + let input = TimestampMicrosecondArray::from(vec![Some(i64::MIN), None]) + .with_timezone_opt(timezone); + let expected = TimestampMicrosecondArray::from(vec![Some(expected), None]) + .with_timezone_opt(timezone); + assert_eq!( + timestamp_trunc(&input, format.into(), true).unwrap(), + expected + ); + } + } + } + + #[test] + fn test_timestamp_trunc_fine_spark42_overflow() { + // Spark 4.2 and later (SPARK-56663) check the SECOND/MILLISECOND subtraction like every + // other unit. First representable boundary for each unit. + for (format, boundary) in [ + ("SECOND", -9_223_372_036_854_000_000), + ("MILLISECOND", -9_223_372_036_854_775_000), + ] { + for timezone in [None, Some("UTC"), Some("Asia/Kolkata")] { + for micros in [i64::MIN, i64::MIN + 1, boundary - 1] { + let input = TimestampMicrosecondArray::from(vec![Some(micros), None]) + .with_timezone_opt(timezone); + assert!(timestamp_trunc(&input, format.into(), false) + .unwrap_err() + .to_string() + .contains("long overflow")); + } + let input = + TimestampMicrosecondArray::from(vec![Some(boundary), Some(boundary + 1), None]) + .with_timezone_opt(timezone); + let expected = + TimestampMicrosecondArray::from(vec![Some(boundary), Some(boundary), None]) + .with_timezone_opt(timezone); + assert_eq!( + timestamp_trunc(&input, format.into(), false).unwrap(), + expected + ); + } + } + } + + #[test] + fn test_timestamp_trunc_dst_gap_and_overlap() { + assert_timestamp_trunc( "HOUR", - "MINUTE", - "SECOND", - "MILLISECOND", - "MICROSECOND", + Some("America/Los_Angeles"), + &[ + Some("2024-11-03T08:30:15.123456Z"), + Some("2024-11-03T09:30:15.123456Z"), + ], + &[Some("2024-11-03T08:00:00Z"), Some("2024-11-03T09:00:00Z")], + ); + assert_timestamp_trunc( + "DAY", + Some("America/New_York"), + &[ + Some("2024-03-10T06:30:15.123456Z"), + Some("2024-03-10T07:30:15.123456Z"), + Some("2024-11-03T06:30:15.123456Z"), + ], + &[ + Some("2024-03-10T05:00:00Z"), + Some("2024-03-10T05:00:00Z"), + Some("2024-11-03T04:00:00Z"), + ], + ); + // Sao Paulo advanced from 23:59:59 on November 3 directly to 01:00 on November 4. + // Spark resolves DAY for a post-gap instant to that day's first valid local time. + assert_timestamp_trunc( + "DAY", + Some("America/Sao_Paulo"), + &[ + Some("2018-11-04T02:30:15.123456Z"), + Some("2018-11-04T03:30:15.123456Z"), + ], + &[Some("2018-11-03T03:00:00Z"), Some("2018-11-04T03:00:00Z")], + ); + } + + #[test] + fn test_timestamp_trunc_coarse_midnight_gap() { + // Toronto's gap starts on the previous date at 23:30 and ends at 00:30. + // atStartOfDay selects 00:30, rather than shifting midnight to 01:00. + for timezone in [ + "America/Toronto", + "Canada/Eastern", + "America/Montreal", + "America/Nassau", ] { - match timestamp_trunc(&array, fmt.to_string()) { - Ok(a) => { - for i in 0..size { - assert!(array.values().get(i) >= a.values().get(i)) - } + assert_timestamp_trunc( + "WEEK", + Some(timezone), + &[ + Some("1919-04-02T16:00:00Z"), + Some("1919-03-31T04:30:00Z"), + None, + ], + &[ + Some("1919-03-31T04:30:00Z"), + Some("1919-03-31T04:30:00Z"), + None, + ], + ); + } + // Asuncion's gap starts exactly at midnight, at both a month and quarter boundary. + for format in ["MONTH", "QUARTER"] { + assert_timestamp_trunc( + format, + Some("America/Asuncion"), + &[Some("2023-10-15T12:00:00Z"), None], + &[Some("2023-10-01T04:00:00Z"), None], + ); + } + } + + #[test] + fn test_timestamp_trunc_coarse_midnight_overlap() { + // Havana repeated local midnight on 2026-11-01. MONTH must use the first occurrence + // (UTC-4), even though the input is after the fall-back transition and has UTC-5. + assert_timestamp_trunc( + "MONTH", + Some("America/Havana"), + &[Some("2026-11-15T12:00:00Z")], + &[Some("2026-11-01T04:00:00Z")], + ); + } + + #[test] + fn test_timestamp_trunc_minute_with_historical_offset_seconds() { + // Africa/Monrovia used UTC-00:44:30 in 1960. The input instant is local 10:30:45, + // which Spark truncates to local 10:30:00 (11:14:30 UTC), not a UTC minute boundary. + assert_timestamp_trunc( + "MINUTE", + Some("Africa/Monrovia"), + &[Some("1960-06-15T11:15:15Z"), None], + &[Some("1960-06-15T11:14:30Z"), None], + ); + + // Asia/Aden changed from UTC+03:06:52 to UTC+03:00 within the local 23:53 minute. + // Truncating the post-transition 23:53:30 must re-resolve local 23:53:00 with the + // pre-transition offset, matching ZonedDateTime.truncatedTo. + assert_timestamp_trunc( + "MINUTE", + Some("Asia/Aden"), + &[Some("1947-03-13T20:53:30.123Z")], + &[Some("1947-03-13T20:46:08Z")], + ); + + // Monrovia's 1972 transition skipped local 00:00:00 through 00:44:29. Truncating + // 00:44:45 targets the gap, which ZonedDateTime shifts forward by 44 minutes 30 seconds. + assert_timestamp_trunc( + "MINUTE", + Some("Africa/Monrovia"), + &[Some("1972-01-07T00:44:45Z")], + &[Some("1972-01-07T01:28:30Z")], + ); + } + + #[test] + fn test_timestamp_trunc_scalar_format_dictionary() { + let mut builder = PrimitiveDictionaryBuilder::::new(); + builder + .append(instant_micros("2024-05-17T12:34:56Z")) + .unwrap(); + builder + .append(instant_micros("2024-06-30T23:59:59Z")) + .unwrap(); + builder + .append(instant_micros("2024-05-17T12:34:56Z")) + .unwrap(); + builder.append_null(); + let input = builder.finish(); + let input = input.with_values(Arc::new( + input + .values() + .as_any() + .downcast_ref::() + .unwrap() + .clone() + .with_timezone_utc(), + )); + let input_keys = input.keys().clone(); + + let result = timestamp_trunc_dyn(&input, "MONTH".to_string(), true).unwrap(); + let result = result + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(result.keys(), &input_keys); + let decoded = result.downcast_dict::().unwrap(); + assert_eq!( + decoded.into_iter().collect::>(), + vec![ + Some(instant_micros("2024-05-01T00:00:00Z")), + Some(instant_micros("2024-06-01T00:00:00Z")), + Some(instant_micros("2024-05-01T00:00:00Z")), + None, + ] + ); + } + + /// Instants around DST and offset transitions, keyed by the session timezone the kernel sees. + /// `None` is a TIMESTAMP_NTZ array. + fn transition_instants() -> Vec<(Option<&'static str>, Vec)> { + let at = |instants: &[&str]| { + instants + .iter() + .map(|i| instant_micros(i)) + .collect::>() + }; + vec![ + // Toronto skipped 1919-03-30 23:30 to 1919-03-31 00:30. + ( + Some("America/Toronto"), + at(&[ + "1919-03-31T04:30:00Z", + "1919-03-31T04:45:00Z", + "1919-04-02T16:00:00Z", + ]), + ), + // Havana repeated midnight on 2020-11-01. + ( + Some("America/Havana"), + at(&[ + "2020-11-01T04:30:00Z", + "2020-11-01T05:30:00Z", + "2020-11-15T12:00:00Z", + ]), + ), + // Sao Paulo skipped midnight on 2018-11-04 and repeated 23:00 on 2019-02-16. + ( + Some("America/Sao_Paulo"), + at(&[ + "2018-11-04T03:30:00Z", + "2018-11-04T12:00:00Z", + "2019-02-17T02:30:00Z", + ]), + ), + ( + Some("America/Los_Angeles"), + at(&[ + "2024-03-10T10:30:00Z", + "2024-03-10T11:15:30Z", + "2024-11-03T08:30:00Z", + "2024-11-03T09:30:00Z", + "1883-06-15T10:30:45Z", + ]), + ), + // Monrovia used -00:44:30 until 1972. + ( + Some("Africa/Monrovia"), + at(&[ + "1960-06-15T11:15:45Z", + "1972-01-07T00:44:29Z", + "1972-01-07T00:44:31Z", + ]), + ), + // Asuncion skipped midnight on 2023-10-01, a MONTH and QUARTER boundary. + ( + Some("America/Asuncion"), + at(&[ + "2023-10-01T03:30:00Z", + "2023-10-01T04:30:00Z", + "2023-10-15T12:00:00Z", + ]), + ), + // Apia skipped 2011-12-30 entirely. + ( + Some("Pacific/Apia"), + at(&[ + "2011-12-30T09:59:59Z", + "2011-12-30T10:00:00Z", + "2011-12-31T12:00:00Z", + ]), + ), + ( + Some("UTC"), + at(&["1500-06-15T12:34:56.123456Z", "3333-05-17T12:34:56.123456Z"]), + ), + ( + None, + vec![ + i64::MAX, + i64::MIN + 400 * MICROS_PER_DAY, + instant_micros("1969-12-31T23:59:59.999999Z"), + instant_micros("2024-05-17T12:34:56.123456Z"), + ], + ), + ] + } + + #[test] + fn row_formats_truncate_like_literal_formats() { + let spellings: Vec<&str> = TIMESTAMP_TRUNC_ALIASES + .iter() + .map(|(name, _)| *name) + .collect(); + for ((timezone, instants), wrap) in transition_instants() + .into_iter() + .flat_map(|zone| [(zone.clone(), true), (zone, false)]) + { + let input = TimestampMicrosecondArray::from(instants).with_timezone_opt(timezone); + let by_spelling: Vec = spellings + .iter() + .map(|spelling| timestamp_trunc(&input, spelling.to_string(), wrap).unwrap()) + .collect(); + // Rotate the spellings across the rows so every row meets every format in a column + // that mixes formats, which exercises the grouping and masking. + for shift in 0..spellings.len() { + let pick = |row: usize| (row + shift) % spellings.len(); + let formats = StringArray::from( + (0..input.len()) + .map(|row| spellings[pick(row)]) + .collect::>(), + ); + let result = timestamp_trunc_array_fmt_dyn(&input, &formats, true).unwrap(); + let result = result + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(result.data_type(), input.data_type()); + for row in 0..input.len() { + assert_eq!( + result.value(row), + by_spelling[pick(row)].value(row), + "{timezone:?} {} at {}", + spellings[pick(row)], + input.value(row) + ); } - _ => unreachable!(), } } } + #[test] + fn row_format_week_resolves_a_midnight_gap_to_its_end() { + // Toronto's gap ran from 1919-03-30 23:30 to 1919-03-31 00:30, so the Monday that WEEK + // lands on starts at 00:30. DAY moves its nonexistent midnight forward by the gap to 01:00. + let input = TimestampMicrosecondArray::from(vec![ + instant_micros("1919-04-02T16:00:00Z"), + instant_micros("1919-03-31T04:45:00Z"), + ]) + .with_timezone("America/Toronto"); + let formats = StringArray::from(vec!["WEEK", "DAY"]); + let result = timestamp_trunc_array_fmt_dyn(&input, &formats, true).unwrap(); + let expected = TimestampMicrosecondArray::from(vec![ + instant_micros("1919-03-31T04:30:00Z"), + instant_micros("1919-03-31T05:00:00Z"), + ]) + .with_timezone("America/Toronto"); + assert_eq!( + result + .as_any() + .downcast_ref::() + .unwrap(), + &expected + ); + } + + #[test] + fn row_formats_handle_nulls_like_spark() { + let input = TimestampMicrosecondArray::from(vec![ + Some(instant_micros("2024-05-17T12:34:56Z")), + Some(instant_micros("2024-05-17T12:34:56Z")), + None, + Some(instant_micros("2024-05-17T12:34:56Z")), + ]) + .with_timezone("America/Los_Angeles"); + // A NULL format gives NULL, and so does a NULL value, whatever its format says. + let formats = StringArray::from(vec![Some("YEAR"), None, Some("not_a_unit"), Some("HOUR")]); + let result = timestamp_trunc_array_fmt_dyn(&input, &formats, true).unwrap(); + let expected = TimestampMicrosecondArray::from(vec![ + Some(instant_micros("2024-01-01T08:00:00Z")), + None, + None, + Some(instant_micros("2024-05-17T12:00:00Z")), + ]) + .with_timezone("America/Los_Angeles"); + assert_eq!( + result + .as_any() + .downcast_ref::() + .unwrap(), + &expected + ); + + let formats = StringArray::from(vec!["YEAR", "not_a_unit", "YEAR", "YEAR"]); + assert!(timestamp_trunc_array_fmt_dyn(&input, &formats, true).is_err()); + } + + #[test] + fn row_format_groups_do_not_fail_each_other() { + // SECOND wraps at the lower bound like Spark before 4.2, while YEAR would overflow there. + // The YEAR group only sees its own row. + let input = + TimestampMicrosecondArray::from(vec![i64::MIN, instant_micros("2024-05-17T12:34:56Z")]); + let formats = StringArray::from(vec!["SECOND", "YEAR"]); + let result = timestamp_trunc_array_fmt_dyn(&input, &formats, true).unwrap(); + let literal = timestamp_trunc(&input.slice(0, 1), "SECOND".to_string(), true).unwrap(); + let result = result + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(result.value(0), literal.value(0)); + assert_eq!(result.value(1), instant_micros("2024-01-01T00:00:00Z")); + + let formats = StringArray::from(vec!["YEAR", "YEAR"]); + assert!(timestamp_trunc_array_fmt_dyn(&input, &formats, true).is_err()); + } + + #[test] + fn row_formats_accept_dictionaries() { + let instants = [ + instant_micros("2024-03-10T10:30:00Z"), + instant_micros("2024-11-03T09:30:00Z"), + instant_micros("2024-11-03T09:30:00Z"), + ]; + let input = + TimestampMicrosecondArray::from(instants.to_vec()).with_timezone("America/Los_Angeles"); + let formats = StringArray::from(vec!["DAY", "HOUR", "MONTH"]); + let expected = timestamp_trunc_array_fmt_dyn(&input, &formats, true).unwrap(); + + let keys = Int32Array::from(vec![0, 1, 1]); + let values = TimestampMicrosecondArray::from(vec![instants[0], instants[1]]) + .with_timezone("America/Los_Angeles"); + let input_dict = DictionaryArray::try_new(keys, Arc::new(values)).unwrap(); + let mut formats_builder = StringDictionaryBuilder::::new(); + for format in ["DAY", "HOUR", "MONTH"] { + formats_builder.append(format).unwrap(); + } + let formats_dict = formats_builder.finish(); + + for (array, formats) in [ + (&input_dict as &dyn Array, &formats as &dyn Array), + (&input as &dyn Array, &formats_dict as &dyn Array), + (&input_dict as &dyn Array, &formats_dict as &dyn Array), + ] { + assert_eq!( + ×tamp_trunc_array_fmt_dyn(array, formats, true).unwrap(), + &expected + ); + } + } + #[test] // test takes too long with miri #[cfg_attr(miri, ignore)] @@ -1179,7 +2839,7 @@ mod tests { } // test cases - if let Ok(a) = timestamp_trunc_array_fmt_dyn(&array, &fmt_array) { + if let Ok(a) = timestamp_trunc_array_fmt_dyn(&array, &fmt_array, true) { for i in 0..array.len() { assert!( array.value(i) @@ -1192,7 +2852,7 @@ mod tests { } else { unreachable!() } - if let Ok(a) = timestamp_trunc_array_fmt_dyn(&array_dict, &fmt_array) { + if let Ok(a) = timestamp_trunc_array_fmt_dyn(&array_dict, &fmt_array, true) { for i in 0..array.len() { assert!( array.value(i) @@ -1205,7 +2865,7 @@ mod tests { } else { unreachable!() } - if let Ok(a) = timestamp_trunc_array_fmt_dyn(&array, &fmt_dict) { + if let Ok(a) = timestamp_trunc_array_fmt_dyn(&array, &fmt_dict, true) { for i in 0..array.len() { assert!( array.value(i) @@ -1218,7 +2878,7 @@ mod tests { } else { unreachable!() } - if let Ok(a) = timestamp_trunc_array_fmt_dyn(&array_dict, &fmt_dict) { + if let Ok(a) = timestamp_trunc_array_fmt_dyn(&array_dict, &fmt_dict, true) { for i in 0..array.len() { assert!( array.value(i) @@ -1233,6 +2893,27 @@ mod tests { } } + #[test] + fn row_formats_follow_the_overflow_policy() { + // Spark before 4.2 wraps SECOND and MILLISECOND below the smallest timestamp, and 4.2 and + // later raise. Cover a column with one format and one that mixes formats. + let input = TimestampMicrosecondArray::from(vec![ + i64::MIN, + i64::MIN, + instant_micros("2024-05-17T12:34:56Z"), + ]) + .with_timezone("UTC"); + for formats in [ + vec!["SECOND", "SECOND", "SECOND"], + vec!["SECOND", "MILLISECOND", "YEAR"], + ] { + let formats = StringArray::from(formats); + assert!(timestamp_trunc_array_fmt_dyn(&input, &formats, true).is_ok()); + let error = timestamp_trunc_array_fmt_dyn(&input, &formats, false).unwrap_err(); + assert!(error.to_string().contains("long overflow"), "{error}"); + } + } + /// Truncating a November timestamp in `America/Denver` to QUARTER must land on the start of /// Q4, which is October 1 — and October 1 is still MDT (UTC-6), not MST (UTC-7). The /// pre-fix kernel reused the input's MST offset for the truncated date, producing a result @@ -1245,7 +2926,7 @@ mod tests { let array = TimestampMicrosecondArray::from(vec![ts_utc_micros]).with_timezone("America/Denver"); - let result = timestamp_trunc(&array, "QUARTER".to_string()).unwrap(); + let result = timestamp_trunc(&array, "QUARTER".to_string(), true).unwrap(); // 2023-10-01 00:00:00 MDT = 2023-10-01 06:00:00 UTC let expected_utc_micros: i64 = 1696140000 * 1_000_000; diff --git a/native/spark-expr/src/math_funcs/internal/decimal_rescale_check.rs b/native/spark-expr/src/math_funcs/internal/decimal_rescale_check.rs index 190d7b060c7..580a01c9afd 100644 --- a/native/spark-expr/src/math_funcs/internal/decimal_rescale_check.rs +++ b/native/spark-expr/src/math_funcs/internal/decimal_rescale_check.rs @@ -195,7 +195,7 @@ impl PhysicalExpr for DecimalRescaleCheckOverflow { Err(_) if self.fail_on_error => spark_cast( arg, &target_type, - &SparkCastOptions::new_without_timezone(EvalMode::Ansi, false), + &SparkCastOptions::new_without_timezone(EvalMode::Ansi), ), Err(error) => Err(error), } diff --git a/native/spark-expr/src/math_funcs/modulo_expr.rs b/native/spark-expr/src/math_funcs/modulo_expr.rs index 5c0652e9a01..09d006c9e49 100644 --- a/native/spark-expr/src/math_funcs/modulo_expr.rs +++ b/native/spark-expr/src/math_funcs/modulo_expr.rs @@ -163,14 +163,14 @@ pub fn create_modulo_expr( let left_256 = Arc::new(Cast::new( left, DataType::Decimal256(p1, s1), - SparkCastOptions::new_without_timezone(EvalMode::Legacy, false), + SparkCastOptions::new_without_timezone(EvalMode::Legacy), None, None, )); let right_256 = Arc::new(Cast::new( right_non_ansi_safe, DataType::Decimal256(p2, s2), - SparkCastOptions::new_without_timezone(EvalMode::Legacy, false), + SparkCastOptions::new_without_timezone(EvalMode::Legacy), None, None, )); @@ -192,7 +192,7 @@ pub fn create_modulo_expr( Ok(Arc::new(Cast::new( modulo_scalar_func, data_type, - SparkCastOptions::new_without_timezone(EvalMode::Legacy, false), + SparkCastOptions::new_without_timezone(EvalMode::Legacy), None, None, ))) diff --git a/native/spark-expr/src/url_funcs/parse_url.rs b/native/spark-expr/src/url_funcs/parse_url.rs index 7e072b68957..ec1b33ab353 100644 --- a/native/spark-expr/src/url_funcs/parse_url.rs +++ b/native/spark-expr/src/url_funcs/parse_url.rs @@ -29,12 +29,14 @@ use arrow::array::{ }; use arrow::datatypes::DataType; use datafusion::common::cast::{as_large_string_array, as_string_array, as_string_view_array}; -use datafusion::common::{exec_datafusion_err, exec_err, Result, ScalarValue}; +use datafusion::common::{exec_err, Result, ScalarValue}; use datafusion::logical_expr::{ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility, }; use regex::Regex; +use crate::SparkError; + // RFC 3986 Appendix B decomposition regex. // Groups: 2=scheme, 4=authority, 5=path, 7=query, 9=fragment static URI_REGEX: LazyLock = LazyLock::new(|| { @@ -109,11 +111,10 @@ fn has_invalid_uri_chars(s: &str) -> bool { } fn invalid_url_err(value: &str) -> datafusion::common::DataFusionError { - exec_datafusion_err!( - "[INVALID_URL] The provided URL '{}' is not valid. Use `try_parse_url` to tolerate \ - malformed URLs and return NULL instead. SQLSTATE: 22P02", - value - ) + SparkError::InvalidUrl { + url: value.to_string(), + } + .into() } fn parse_url_component(value: &str, part: &str, key: Option<&str>) -> Result> { diff --git a/spark/src/main/java/org/apache/comet/NativeBase.java b/spark/src/main/java/org/apache/comet/NativeBase.java index ab4eee856ba..7952e736116 100644 --- a/spark/src/main/java/org/apache/comet/NativeBase.java +++ b/spark/src/main/java/org/apache/comet/NativeBase.java @@ -384,4 +384,16 @@ public static void releaseNative() throws Throwable { * @return the version compiled into libcomet */ public static native String getTzdataVersion(); + + /** + * The comma-joined URL schemes the native Iceberg storage factory can open, for reads ({@code + * forWrite} false) or writes ({@code forWrite} true). This is the authoritative list the JVM + * Iceberg scan and write gates load, so the planner never hardcodes (and drifts from) the set the + * native factory actually builds. Opt-in S3-compliant alias schemes are not included; the JVM + * adds them from {@code fs.comet.s3Compliant.schemes}. + * + * @param forWrite true for the write path's schemes, false for the scan path's + * @return the supported schemes joined by commas, e.g. "file,memory,gs,s3,s3a" + */ + public static native String icebergStorageSchemes(boolean forWrite); } diff --git a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java index d2881bfe449..3ab34fe0496 100644 --- a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java +++ b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java @@ -32,6 +32,8 @@ import org.apache.spark.comet.CometTaskContextShim; import org.apache.spark.util.TaskCompletionListener; +import com.google.common.annotations.VisibleForTesting; + import org.apache.comet.util.ClassLoaders; /** @@ -86,6 +88,8 @@ public class CometUdfBridge { * @param numRows row count of the current batch. Mirrors DataFusion's {@code * ScalarFunctionArgs.number_rows}; the only batch-size signal a zero-input UDF (e.g. a * zero-arg non-deterministic ScalaUDF) ever sees. + * @param partitionIndex index of the partition the calling native plan computes + * @param planId id of the calling native plan * @param taskContext propagated Spark {@link TaskContext} from the driving Spark task thread, or * {@code null} outside a Spark task. Treated as ground truth for the call: installed as the * thread-local on entry, with the prior value (if any) saved and restored in {@code finally}. @@ -113,6 +117,8 @@ public static void evaluate( long outArrayPtr, long outSchemaPtr, int numRows, + int partitionIndex, + long planId, TaskContext taskContext, ClassLoader classLoader) { assert udfClassName != null && !udfClassName.isEmpty() : "udfClassName must be non-empty"; @@ -148,6 +154,8 @@ public static void evaluate( outArrayPtr, outSchemaPtr, numRows, + partitionIndex, + planId, taskContext); } finally { // Unconditional: a no-op when nothing was installed, and it also undoes any change the @@ -170,6 +178,8 @@ private static void evaluateInternal( long outArrayPtr, long outSchemaPtr, int numRows, + int partitionIndex, + long planId, TaskContext taskContext) { long taskAttemptId = (taskContext != null) ? taskContext.taskAttemptId() : NO_TASK_ID; @@ -228,7 +238,7 @@ private static void evaluateInternal( inputs[i] = Data.importVector(importAllocator, inArr, inSch, null); } - result = udf.evaluate(inputs, numRows); + result = udf.evaluate(inputs, numRows, partitionIndex, planId); // Recorded before the checks below, so the invariant holds however this exits. resultIsInput = isOneOf(result, inputs); if (!(result instanceof FieldVector)) { @@ -290,6 +300,27 @@ private static void evaluateInternal( } } + /** + * Called by {@code CometExecIterator} once native plan {@code planId} of task {@code + * taskAttemptId} has closed. Passes the plan to every {@link CometUDF} instance of the task, see + * {@code CometUDF.releasePlan}. + */ + public static void releasePlan(long taskAttemptId, long planId) { + ConcurrentHashMap perTask = INSTANCES.get(taskAttemptId); + if (perTask != null) { + for (CometUDF udf : perTask.values()) { + udf.releasePlan(planId); + } + } + } + + /** The instance of {@code udfClassName} that task {@code taskAttemptId} holds, or null. */ + @VisibleForTesting + public static CometUDF instanceFor(long taskAttemptId, String udfClassName) { + ConcurrentHashMap perTask = INSTANCES.get(taskAttemptId); + return perTask == null ? null : perTask.get(udfClassName); + } + /** Whether the UDF handed back one of the vectors it was given, rather than a new one. */ private static boolean isOneOf(ValueVector result, ValueVector[] inputs) { for (ValueVector input : inputs) { diff --git a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java index afdaafd828b..a468446bae7 100644 --- a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java +++ b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java @@ -20,6 +20,7 @@ package org.apache.spark; import java.io.IOException; +import java.util.NoSuchElementException; import java.util.concurrent.atomic.AtomicLong; import org.slf4j.Logger; @@ -38,6 +39,9 @@ public class CometTaskMemoryManager { private static final Logger logger = LoggerFactory.getLogger(CometTaskMemoryManager.class); + /** Attempts at an acquire whose task entry Spark keeps losing, see acquireFromSpark. */ + private static final int MAX_ACQUIRE_ATTEMPTS = 3; + /** The id of the native plan this memory manager was created for, the first in its task. */ private final long id; @@ -70,7 +74,7 @@ public long acquireMemory(long size) { if (logger.isTraceEnabled()) { logger.trace("Task {} requested {} bytes", taskAttemptId, size); } - long acquired = internal.acquireExecutionMemory(size, nativeMemoryConsumer); + long acquired = acquireFromSpark(size); long newUsed = used.addAndGet(acquired); // A partial grant is routine, not an error: the native pool either refuses the reservation, // which tells the operator to spill, or carries the shortfall as overcommit. A refusal that @@ -94,6 +98,57 @@ public long acquireMemory(long size) { return acquired; } + /** + * Acquires from Spark, retrying when an acquire waiting in Spark lost the task's entry. Spark + * removes a task's entry from its execution pool when the task's balance reaches zero, and an + * acquire that was waiting for memory then fails when it wakes up with a NoSuchElementException + * ("key not found: " and the task id). Releases take no task monitor, so another native thread or + * a JVM consumer of the same task can empty the balance while an acquire waits. The failed call + * was granted nothing, because Spark only drops the entry at a zero balance, which cannot hold + * while this call has a partial grant. So trying again is safe: it registers the task again and + * waits for its share as Spark would have. After a few failed attempts the grant is refused. The + * helper holds no lock of its own across attempts, so a retry parks in Spark exactly as the first + * attempt did. + */ + private long acquireFromSpark(long size) { + for (int attempt = 1; ; attempt++) { + try { + return internal.acquireExecutionMemory(size, nativeMemoryConsumer); + } catch (NoSuchElementException e) { + if (!isMissingTaskEntry(e)) { + throw e; + } + if (attempt >= MAX_ACQUIRE_ATTEMPTS) { + logger.warn( + "Task {} lost its execution memory entry in Spark on {} attempts to acquire {} " + + "bytes, refusing the request", + taskAttemptId, + attempt, + size); + return 0; + } + logger.info( + "Task {} lost its execution memory entry in Spark while waiting to acquire {} bytes, " + + "trying again", + taskAttemptId, + size); + } + } + } + + /** Whether {@code e} is Spark's execution pool failing to find this task's entry. */ + private boolean isMissingTaskEntry(NoSuchElementException e) { + if (!("key not found: " + taskAttemptId).equals(e.getMessage())) { + return false; + } + for (StackTraceElement frame : e.getStackTrace()) { + if ("org.apache.spark.memory.ExecutionMemoryPool".equals(frame.getClassName())) { + return true; + } + } + return false; + } + // Called by Comet native through JNI public void releaseMemory(long size) { if (logger.isTraceEnabled()) { diff --git a/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java b/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java index b7c7c58848e..e3a2269cacc 100644 --- a/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java +++ b/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java @@ -20,10 +20,18 @@ package org.apache.spark.shuffle.comet; import java.io.IOException; +import java.util.Map; +import java.util.NoSuchElementException; +import java.util.function.Supplier; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import org.apache.spark.memory.MemoryConsumer; import org.apache.spark.memory.MemoryMode; +import org.apache.spark.memory.SparkOutOfMemoryError; import org.apache.spark.memory.TaskMemoryManager; +import org.apache.spark.unsafe.array.LongArray; import org.apache.spark.unsafe.memory.MemoryBlock; /** @@ -38,6 +46,12 @@ */ public final class CometUnifiedShuffleMemoryAllocator extends CometShuffleMemoryAllocatorTrait { + private static final Logger logger = + LoggerFactory.getLogger(CometUnifiedShuffleMemoryAllocator.class); + + /** Attempts at an allocation whose task entry Spark keeps losing, see retryMissingTaskEntry. */ + private static final int MAX_ALLOCATE_ATTEMPTS = 3; + CometUnifiedShuffleMemoryAllocator(TaskMemoryManager taskMemoryManager, long pageSize) { super(taskMemoryManager, pageSize, MemoryMode.OFF_HEAP); if (taskMemoryManager.getTungstenMemoryMode() != MemoryMode.OFF_HEAP) { @@ -54,7 +68,78 @@ public long spill(long l, MemoryConsumer memoryConsumer) throws IOException { } public synchronized MemoryBlock allocate(long required) { - return this.allocatePage(required); + return retryMissingTaskEntry(required, () -> this.allocatePage(required)); + } + + /** + * Allocates a pointer array, retrying like {@link #allocate}. The inherited method asks Spark for + * the page directly rather than through allocate, and the sorters' pointer arrays come through + * here. + */ + @Override + public LongArray allocateArray(long size) { + return retryMissingTaskEntry(size * 8L, () -> super.allocateArray(size)); + } + + /** + * Runs an allocation through Spark, retrying when an allocation waiting in Spark lost the task's + * entry. Spark removes a task's entry from its execution pool when the task's balance reaches + * zero, and an allocation that was waiting for memory then fails when it wakes up with a + * NoSuchElementException ("key not found: " and the task id). Releases take no task monitor, so a + * native consumer or another JVM consumer of the same task can empty the balance while an + * allocation waits. The failed call was granted nothing and took no page, because Spark only + * drops the entry at a zero balance, so trying again is safe: it registers the task again and + * waits for its share as Spark would have. After a few failed attempts the allocation fails with + * a SparkOutOfMemoryError, which the shuffle writers handle like any other refused page, caused + * by the last attempt's exception. The Spark issue is SPARK-59444. + */ + private T retryMissingTaskEntry(long required, Supplier allocation) { + for (int attempt = 1; ; attempt++) { + try { + return allocation.get(); + } catch (NoSuchElementException e) { + if (!isMissingTaskEntry(e)) { + throw e; + } + if (attempt >= MAX_ALLOCATE_ATTEMPTS) { + logger.warn( + "Lost the task's execution memory entry in Spark on {} attempts to allocate {} " + + "bytes, refusing the allocation", + attempt, + required, + e); + SparkOutOfMemoryError error = + new SparkOutOfMemoryError( + "UNABLE_TO_ACQUIRE_MEMORY", + Map.of( + "requestedBytes", String.valueOf(required), + "receivedBytes", String.valueOf(0))); + error.initCause(e); + throw error; + } + logger.info( + "Lost the task's execution memory entry in Spark while waiting to allocate {} bytes, " + + "trying again", + required); + } + } + } + + /** + * Whether {@code e} is Spark's execution pool failing to find the task's entry. The task id is + * not available here, so this matches the message prefix and the pool's stack frame. + */ + private static boolean isMissingTaskEntry(NoSuchElementException e) { + String message = e.getMessage(); + if (message == null || !message.startsWith("key not found: ")) { + return false; + } + for (StackTraceElement frame : e.getStackTrace()) { + if ("org.apache.spark.memory.ExecutionMemoryPool".equals(frame.getClassName())) { + return true; + } + } + return false; } public synchronized long free(MemoryBlock block) { diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 09167333deb..4aab7c72a11 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -178,10 +178,12 @@ object CometConf extends ShimCometConf { "When enabled, the native Parquet reader evaluates pushed filters during decode " + "and lazily materializes projected columns for surviving rows (DataFusion's " + "pushdown_filters / late-materialization). Format-level pruning (row-group " + - "statistics, page index, bloom filters) is independent of this flag and runs " + - "whenever Spark's spark.sql.parquet.filterPushdown is enabled. Disabling this " + - "flag still lets format-level pruning work; the per-row eval falls back to " + - "the CometFilter operator above the scan.") + "statistics, page index, bloom filters) runs whenever Spark's " + + "spark.sql.parquet.filterPushdown is enabled, except that while this flag is " + + "enabled, a pushed filter that compares FLOAT or DOUBLE values does not prune, " + + "because the reader must evaluate it with Spark's NaN and signed-zero semantics. " + + "Disabling this flag still lets format-level pruning work; the per-row eval falls " + + "back to the CometFilter operator above the scan.") .booleanConf .createWithDefault(false) @@ -244,9 +246,10 @@ object CometConf extends ShimCometConf { .category(CATEGORY_EXEC) .doc( "When enabled, the single row that a query without a FROM clause, such as `SELECT 1`, " + - "reads will be converted to Arrow format.") + "reads will be converted to Arrow format. The row has no columns, so converting it " + + "costs almost nothing, and it lets the operators above it run in Comet.") .booleanConf - .createWithDefault(false) + .createWithDefault(true) val COMET_CONVERT_FROM_ROW_DATA_SOURCE_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.convert.rowDataSource.enabled") @@ -270,6 +273,20 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + val COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.convert.shuffleInput.enabled") + .category(CATEGORY_EXEC) + .doc( + "When enabled, a shuffle that would use Comet's JVM columnar shuffle because its input " + + "comes from a Spark operator converts that input to Arrow format and uses Comet " + + "native shuffle instead. This applies only where native shuffle supports the " + + "partitioning and Comet can convert all of the input's columns, which rules out " + + "calendar intervals. A shuffle that hashes a decimal wider than 18 digits, or a " + + "string or a value computed from a string, stays on the JVM columnar shuffle, so " + + "that it partitions rows as Spark does.") + .booleanConf + .createWithDefault(false) + val COMET_EXEC_ENABLED: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.enabled") .category(CATEGORY_EXEC) .doc( @@ -358,9 +375,11 @@ object CometConf extends ShimCometConf { .category(CATEGORY_EXEC) .doc( "Whether to enable Comet native scans and fused Spark reads of in-memory cached tables. " + - "Requires spark.comet.enabled=true. At startup, this setting also decides whether " + - "CometDriverPlugin installs Comet's cache serializer, which stores cached data in " + - "Arrow format. The plugin installs it only if spark.comet.enabled and " + + "Requires spark.comet.enabled=true. Defaults to true from Spark 3.5 and to false on " + + "Spark 3.4, where AQE cannot coalesce the shuffle partitions of a union that Comet " + + "runs with a cached relation in one branch. At startup, this setting also decides " + + "whether CometDriverPlugin installs Comet's cache serializer, which stores cached " + + "data in Arrow format. The plugin installs it only if spark.comet.enabled and " + "spark.comet.exec.enabled are also enabled at startup, and only with one of Comet's " + "shuffle managers while Comet shuffle is enabled. " + "Because spark.sql.cache.serializer is a " + @@ -381,7 +400,9 @@ object CometConf extends ShimCometConf { "soon as it is serialized, including the disk half of the default " + "MEMORY_AND_DISK storage level.") .booleanConf - .createWithDefault(false) + // CometCoalesceShufflePartitions needs a query-stage optimizer rule hook, which Spark 3.4 + // lacks. + .createWithDefault(CometSparkSessionExtensions.isSpark35Plus) val COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_CODEC: ConfigEntry[String] = conf("spark.comet.exec.inMemoryCache.compression.codec") @@ -542,6 +563,22 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + val COMET_FORCE_SHJ_MAX_BUILD_SIZE: OptionalConfigEntry[Long] = + conf(s"$COMET_EXEC_CONFIG_PREFIX.forceShuffledHashJoin.maxBuildSize") + .category(CATEGORY_EXEC) + .doc(s"The build side size below which `${COMET_FORCE_SHJ.key}` converts a " + + "SortMergeJoin to ShuffledHashJoin. The size is Spark's planning estimate of the build " + + "child, or under AQE the materialized shuffle size of the build side. A build side at " + + "or over this size, or one with no statistics, keeps the SortMergeJoin. When unset, " + + "Spark's own rule applies: `spark.sql.autoBroadcastJoinThreshold` times the initial " + + "shuffle partition count (`spark.sql.adaptive.coalescePartitions.initialPartitionNum` " + + "when AQE and partition coalescing are both on and it is set, else " + + "`spark.sql.shuffle.partitions`). When broadcasts are disabled with a non-positive " + + "threshold, Spark's default threshold of 10 MB is used instead. A " + + s"non-positive value removes the limit. $TUNING_GUIDE.") + .bytesConf(ByteUnit.BYTE) + .createOptional + val COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.aggregate.skipPartial.enabled") .category(CATEGORY_EXEC) @@ -725,14 +762,14 @@ object CometConf extends ShimCometConf { conf("spark.comet.shuffle.revertRedundantColumnar.enabled") .withAlternative(s"$COMET_EXEC_CONFIG_PREFIX.shuffle.revertRedundantColumnar.enabled") .category(CATEGORY_SHUFFLE) - .doc( - "When enabled, Comet reverts a `CometShuffleExchangeExec` with `CometColumnarShuffle` " + - "back to Spark's `ShuffleExchangeExec` when both its parent and child are non-Comet " + - "hash aggregate operators. This avoids a redundant " + - "row -> Arrow -> shuffle -> Arrow -> row conversion when no Comet operator on either " + - "side can consume columnar output. Disable to keep Comet columnar shuffle even in " + - "that case, which preserves Comet's off-heap shuffle memory accounting at the cost of " + - "the extra conversion.") + .doc("When enabled, Comet reverts a `CometShuffleExchangeExec` with `CometColumnarShuffle` " + + "back to Spark's `ShuffleExchangeExec` when both its parent and child are non-Comet " + + "aggregate operators. The same applies to a native shuffle whose input " + + s"`${COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.key}` converted. This avoids a redundant " + + "row -> Arrow -> shuffle -> Arrow -> row conversion when no Comet operator on either " + + "side can consume columnar output. Disable to keep the Comet shuffle even in " + + "that case, which preserves Comet's off-heap shuffle memory accounting at the cost of " + + "the extra conversion.") .booleanConf .createWithDefault(true) @@ -1190,9 +1227,7 @@ object CometConf extends ShimCometConf { "When enabled, fall back to Spark for floating-point operations that may differ from " + "Spark, such as comparing -0.0 and 0.0. `ORDER BY`, window ordering and range " + "partitioning keys are unaffected, including floating-point values nested in arrays " + - "and structs, because Comet normalizes those comparison keys to match Spark. The " + - "exception is a nested key whose type can hold a null element or field, which falls " + - "back until the native sort and window frames order those nulls as Spark does. " + + "and structs, because Comet normalizes those comparison keys to match Spark. " + "`sort_array` is unaffected too, because it follows Spark's ordering. " + s"$COMPAT_GUIDE.") .booleanConf @@ -1264,6 +1299,19 @@ object CometConf extends ShimCometConf { .bytesConf(ByteUnit.BYTE) .createWithDefault(100L * 1024 * 1024 * 1024) // 100 GB + // Used on native side. Check spark_config.rs how the config is used + val COMET_EXEC_SPILL_COMPRESSION_CODEC: ConfigEntry[String] = + conf(s"$COMET_EXEC_CONFIG_PREFIX.spill.compression.codec") + .category(CATEGORY_TUNING) + .doc( + "The codec used to compress the files that native sorts, aggregations, and joins " + + "spill to disk. lz4 and zstd are supported, and none disables compression. zstd " + + "writes smaller files than lz4 but uses more CPU. Comet's shuffle writers use " + + s"${COMET_SHUFFLE_COMPRESSION_CODEC.key} instead. $TUNING_GUIDE.") + .stringConf + .checkValues(Set("lz4", "zstd", "none")) + .createWithDefault("lz4") + val COMET_RESPECT_DATAFUSION_CONFIGS: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.respectDataFusionConfigs") .category(CATEGORY_TESTING) diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 22143b4f188..27f880f9132 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -26,6 +26,7 @@ import java.util.concurrent.atomic.AtomicBoolean import scala.util.control.NonFatal +import org.apache.arrow.c.ArrowArrayStream import org.apache.arrow.memory.BufferAllocator import org.apache.hadoop.conf.Configuration import org.apache.spark._ @@ -33,6 +34,7 @@ import org.apache.spark.broadcast.Broadcast import org.apache.spark.internal.Logging import org.apache.spark.network.util.ByteUnit import org.apache.spark.sql.comet.CometMetricNode +import org.apache.spark.sql.comet.execution.arrow.CometArrowStream import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.vectorized._ import org.apache.spark.util.SerializableConfiguration @@ -43,6 +45,7 @@ import org.apache.comet.exceptions.CometQueryExecutionException import org.apache.comet.parquet.CometFileKeyUnwrapper import org.apache.comet.serde.Config.ConfigMap import org.apache.comet.shuffle.ShufflePartitionPusher +import org.apache.comet.udf.CometUdfBridge import org.apache.comet.vector.NativeUtil /** @@ -224,6 +227,17 @@ class CometExecIterator( CometExecIterator.startMemoryUsageLog() + /** + * Matches any throwable once a JVM input of this plan has failed, and extracts what that input + * threw. Native saw only its text. See `CometArrowStream.inputFailure`. + */ + private object InputFailure { + def unapply(nativeFailure: Throwable): Option[Throwable] = + inputObjects.iterator + .collect { case stream: ArrowArrayStream => CometArrowStream.inputFailure(stream) } + .collectFirst { case Some(failure) => failure } + } + private def getNextBatch: Option[ColumnarBatch] = { assert(partitionIndex >= 0 && partitionIndex < numParts) @@ -246,6 +260,10 @@ class CometExecIterator( result } catch { + // Native saw only the text of what a JVM input threw, so rethrow the exception itself. + case InputFailure(failure) => + throw failure + // Handle CometQueryExecutionException with JSON payload first case e: CometQueryExecutionException => logError(s"Native execution for task $taskAttemptId failed", e) @@ -339,6 +357,10 @@ class CometExecIterator( // global refs. attempt(nativeLib.releasePlan(plan)) + // Drops what the task's JVM UDFs keep for this plan, such as its dispatched kernels. `id` is + // the plan id the native plan passed them. + attempt(CometUdfBridge.releasePlan(taskAttemptId, id)) + // Run the diagnostics even when teardown failed: a failed teardown is exactly when the // non-zero memory usage warning below is most informative. attempt { @@ -853,6 +875,7 @@ object CometExecIterator extends Logging { CometConf.COMET_DEBUG_ENABLED, CometConf.COMET_DEBUG_MEMORY_ENABLED, CometConf.COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED, + CometConf.COMET_EXEC_SPILL_COMPRESSION_CODEC, CometConf.COMET_EXPLAIN_NATIVE_ENABLED, CometConf.COMET_MAX_TEMP_DIRECTORY_SIZE, CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala index 2052539efec..e2edb3c989c 100644 --- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala +++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala @@ -24,8 +24,12 @@ import org.apache.arrow.vector._ import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector} import org.apache.arrow.vector.types.pojo.Field import org.apache.spark.internal.Logging -import org.apache.spark.sql.catalyst.expressions.{BoundReference, Expression, Literal, Unevaluable} +import org.apache.spark.sql.catalyst.DeserializerBuildHelper.createDeserializerForTypesSupportValueOf +import org.apache.spark.sql.catalyst.SerializerBuildHelper.{createSerializerForBoolean, createSerializerForByte, createSerializerForDouble, createSerializerForFloat, createSerializerForInteger, createSerializerForLong, createSerializerForShort} +import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder +import org.apache.spark.sql.catalyst.expressions.{BoundReference, Expression, Literal, ScalaUDF, Unevaluable} import org.apache.spark.sql.catalyst.expressions.codegen._ +import org.apache.spark.sql.execution.ExecSubqueryExpression import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ @@ -46,7 +50,8 @@ import org.apache.comet.shims.{CometExprTraitShim, CometTypeShim} * Input- and output-side emission live in [[CometBatchKernelCodegenInput]] and * [[CometBatchKernelCodegenOutput]]. This file owns the [[ArrowColumnSpec]] vocabulary, the * [[canHandle]] / [[allocateOutput]] / [[compile]] / [[generateSource]] entry points, and - * cross-cutting kernel-shape decisions (NullIntolerant short-circuit, CSE variant). + * cross-cutting kernel-shape decisions (NullIntolerant short-circuit, CSE variant, boxed + * primitive `ScalaUDF` values without their encoders). * * The generated kernel is the `InternalRow` that Spark's `BoundReference.genCode` reads from. See * [[generateSource]] for how the wiring is set up. @@ -173,11 +178,9 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come // instance with a single `init(partitionIndex)` call, so `Rand` / `MonotonicallyIncreasingID` // state advances correctly across batches. // - // `ExecSubqueryExpression` (`ScalarSubquery`, `InSubqueryExec`) is accepted: the surrounding - // Comet operator's inherited `SparkPlan.waitForSubqueries` populates the subquery's - // `result` field before evaluation. The closure serializer captures that value into the - // arg-0 bytes, and the dispatcher keys its compile cache on those bytes, so distinct subquery - // results produce distinct cache entries. + // `ExecSubqueryExpression` (`ScalarSubquery`, `InSubqueryExec`): rejected. The tree is + // closure-serialized at plan time, before Spark has run the subquery, so the deserialized + // copy never holds a result and `ScalarSubquery.doGenCode` fails with "has not finished". // // `Unevaluable`: rejected by default. `isCodegenInertUnevaluable` exempts version-specific // leaves that are `Unevaluable` but never invoked by codegen (e.g. Spark 4.0's @@ -188,6 +191,7 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come boundExpr.find { case _: org.apache.spark.sql.catalyst.expressions.aggregate.AggregateFunction => true case _: org.apache.spark.sql.catalyst.expressions.Generator => true + case _: ExecSubqueryExpression => true case u: Unevaluable if isCodegenInertUnevaluable(u) => false case _: Unevaluable => true case _: org.apache.comet.udf.NativeUdfCall => true @@ -196,7 +200,7 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come case Some(bad) => return Some( s"codegen dispatch: expression ${bad.getClass.getSimpleName} not supported " + - "(aggregate, generator, unevaluable, or native UDF)") + "(aggregate, generator, subquery, unevaluable, or native UDF)") case None => } val badRef = boundExpr.collectFirst { @@ -263,6 +267,7 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come inputSchema: Seq[ArrowColumnSpec]): GeneratedSource = { canHandle(boundExpr).foreach(reason => throw new IllegalArgumentException(s"CometBatchKernelCodegen.generateSource: $reason")) + val expr = withoutBoxedPrimitiveEncoders(boundExpr) val ctx = new CodegenContext // `BoundReference.genCode` emits `${ctx.INPUT_ROW}.getUTF8String(ord)`. Aliasing `row` to // `this` at the top of `process` routes those reads to the kernel's typed getters (final @@ -293,19 +298,19 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come // `subexprFunctionsCode` is the concatenated helper invocation block, spliced into the // per-row body by `defaultBody`. val ev = if (SQLConf.get.subexpressionEliminationEnabled) { - ctx.generateExpressions(Seq(boundExpr), doSubexpressionElimination = true).head + ctx.generateExpressions(Seq(expr), doSubexpressionElimination = true).head } else { - boundExpr.genCode(ctx) + expr.genCode(ctx) } val subExprsCode = ctx.subexprFunctionsCode val (cls, setup, snippet) = - CometBatchKernelCodegenOutput.emitOutputWriter(boundExpr.dataType, ev.value, ctx) - (cls, setup, defaultBody(boundExpr, inputSchema, ev, snippet, subExprsCode)) + CometBatchKernelCodegenOutput.emitOutputWriter(expr.dataType, ev.value, ctx) + (cls, setup, defaultBody(expr, inputSchema, ev, snippet, subExprsCode)) } val typedFieldDecls = CometBatchKernelCodegenInput.emitInputFieldDecls(inputSchema) val typedInputCasts = CometBatchKernelCodegenInput.emitInputCasts(inputSchema) - val decimalTypeByOrdinal = CometBatchKernelCodegenInput.decimalPrecisionByOrdinal(boundExpr) + val decimalTypeByOrdinal = CometBatchKernelCodegenInput.decimalPrecisionByOrdinal(expr) val getters = CometBatchKernelCodegenInput.emitTypedGetters(inputSchema, decimalTypeByOrdinal) val nested = CometBatchKernelCodegenInput.emitNestedClasses(inputSchema) @@ -371,6 +376,73 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come GeneratedSource(code.body, code, ctx.references.toArray) } + /** + * Drops the encoder from each boxed primitive parameter and result of every `ScalaUDF` in the + * tree, so the kernel passes those values to and from the user function directly (#6706). + * + * A typed Scala UDF carries an `ExpressionEncoder` for each parameter and one for its result. + * For a boxed primitive (`(x: java.lang.Long) => ...`) the parameter's deserializer is + * `java.lang.Long.valueOf(input[0])`, which yields null for a null input, and the result's + * serializer is `input[0].longValue()`. `ScalaUDF` runs each one on every row through a + * projection: a `SafeProjection` into a `GenericInternalRow` going in, and an + * `UnsafeProjection` into an `UnsafeRow` coming out. Spark's whole-stage codegen runs the same + * projections. + * + * Without an encoder, `ScalaUDF` converts the value with `CatalystTypeConverters` instead, + * which passes a primitive through in its boxed form: `identity` going in, and coming out an + * `Option` unwrap that a boxed value never takes. The function receives the same value or null, + * and Catalyst gets the same result. Spark 4 builds a Java UDF the same way, with no encoders. + * + * Only an encoder whose expression is exactly the one Spark builds for a boxed primitive is + * dropped. Every other encoder (`String`, `Option`, case classes, collections) stays. + */ + private[comet] def withoutBoxedPrimitiveEncoders(expr: Expression): Expression = + expr.transformUp { case udf: ScalaUDF => + udf.copy( + inputEncoders = udf.inputEncoders.zip(udf.children).map { case (enc, child) => + enc.filterNot(isBoxedPrimitiveDeserializer(_, child.dataType)) + }, + outputEncoder = udf.outputEncoder.filterNot(isBoxedPrimitiveSerializer(_, udf.dataType))) + } + + /** + * Boxed class of each type `CatalystTypeConverters` passes through, and Spark's builder for the + * serializer of that class's encoder. + */ + private def boxedPrimitive(dt: DataType): Option[(Class[_], Expression => Expression)] = + dt match { + case BooleanType => Some((classOf[java.lang.Boolean], createSerializerForBoolean)) + case ByteType => Some((classOf[java.lang.Byte], createSerializerForByte)) + case ShortType => Some((classOf[java.lang.Short], createSerializerForShort)) + case IntegerType => Some((classOf[java.lang.Integer], createSerializerForInteger)) + case LongType => Some((classOf[java.lang.Long], createSerializerForLong)) + case FloatType => Some((classOf[java.lang.Float], createSerializerForFloat)) + case DoubleType => Some((classOf[java.lang.Double], createSerializerForDouble)) + case _ => None + } + + /** + * True iff `enc` deserializes a `dt` value exactly as the encoder of its boxed class does, with + * `valueOf` over a nullable reference. The reference must be nullable because `StaticInvoke` + * propagates null only from a nullable argument: over a non-nullable one, Spark would box the 0 + * that a row holds for a null. + */ + private def isBoxedPrimitiveDeserializer(enc: ExpressionEncoder[_], dt: DataType): Boolean = + boxedPrimitive(dt).exists { case (boxed, _) => + enc.objDeserializer == + createDeserializerForTypesSupportValueOf(BoundReference(0, dt, nullable = true), boxed) + } + + /** + * True iff `enc` serializes a boxed `dt` value exactly as the encoder of the boxed class does, + * by unboxing it. `ScalaUDF` turns a null result into null before it calls the serializer, so + * the serializer only ever sees a boxed value. + */ + private def isBoxedPrimitiveSerializer(enc: ExpressionEncoder[_], dt: DataType): Boolean = + boxedPrimitive(dt).exists { case (boxed, serializerFor) => + enc.objSerializer == serializerFor(BoundReference(0, ObjectType(boxed), nullable = true)) + } + /** * Per-row body. For `NullIntolerant` expressions whose input nulls fully determine a null * output (see [[canShortCircuitNulls]]), prepends a short-circuit on those ordinals so the diff --git a/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala b/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala index e14fb086b12..b6fcdf32806 100644 --- a/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala +++ b/spark/src/main/scala/org/apache/comet/expressions/CometCast.scala @@ -23,7 +23,6 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, Cast, Expression, L import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, DataType, DataTypes, DecimalType, MapType, NullType, StructType, TimestampNTZType, TimestampType} -import org.apache.comet.CometConf import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, withFallbackReason} import org.apache.comet.DataTypeSupport.isComplexType import org.apache.comet.serde.{CodegenDispatchFallback, CometExpressionSerde, CometTimeZone, Compatible, ExprOuterClass, Incompatible, SupportLevel, Unsupported} @@ -167,10 +166,6 @@ object CometCast castBuilder.setChild(childExpr) castBuilder.setDatatype(dataType) castBuilder.setEvalMode(evalModeToProto(evalMode)) - castBuilder.setAllowIncompat( - SQLConf.get - .getConfString(CometConf.getExprAllowIncompatConfigKey(classOf[Cast]), "false") - .toBoolean) castBuilder.setTimezone(timeZone) castBuilder.setIsSpark4Plus(isSpark40Plus) Some( diff --git a/spark/src/main/scala/org/apache/comet/expressions/CometEvalMode.scala b/spark/src/main/scala/org/apache/comet/expressions/CometEvalMode.scala index 59e9c89a6b5..d61f0eb3c7b 100644 --- a/spark/src/main/scala/org/apache/comet/expressions/CometEvalMode.scala +++ b/spark/src/main/scala/org/apache/comet/expressions/CometEvalMode.scala @@ -31,12 +31,4 @@ package org.apache.comet.expressions */ object CometEvalMode extends Enumeration { val LEGACY, ANSI, TRY = Value - - def fromBoolean(ansiEnabled: Boolean): Value = if (ansiEnabled) { - ANSI - } else { - LEGACY - } - - def fromString(str: String): CometEvalMode.Value = CometEvalMode.withName(str) } diff --git a/spark/src/main/scala/org/apache/comet/iceberg/DeltaLogicalFields.scala b/spark/src/main/scala/org/apache/comet/iceberg/DeltaLogicalFields.scala new file mode 100644 index 00000000000..feebad3a100 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/iceberg/DeltaLogicalFields.scala @@ -0,0 +1,46 @@ +/* + * 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.iceberg + +import org.apache.spark.sql.catalyst.analysis.NamedRelation +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.catalyst.util.WriteDeltaProjections +import org.apache.spark.sql.connector.write.Write + +sealed trait DeltaCommand extends Product with Serializable +case object DeltaDelete extends DeltaCommand +case object DeltaUpdate extends DeltaCommand +case object DeltaMerge extends DeltaCommand + +case class DeltaLogicalFields( + query: LogicalPlan, + originalTable: NamedRelation, + projections: WriteDeltaProjections, + write: Option[Write], + command: Option[DeltaCommand] = None) + +private[iceberg] trait IcebergDeltaLogicalShimApi { + def extract(plan: LogicalPlan): Option[DeltaLogicalFields] +} + +/** Version-specific extraction of Spark's public WriteDelta logical node. */ +private[iceberg] trait IcebergDeltaLogicalFieldsShimApi { + def extract(plan: LogicalPlan): Option[DeltaLogicalFields] +} diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala index 9a0dfee85fb..96626631393 100644 --- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala @@ -205,6 +205,20 @@ object IcebergReflection extends Logging { batchWrite.getClass.getName.startsWith(ClassNames.SPARK_WRITE + "$") } + private val positionDeltaWriteClassName = + "org.apache.iceberg.spark.source.SparkPositionDeltaWrite" + private val positionDeltaBatchWriteClassPrefix = + "org.apache.iceberg.spark.source.SparkPositionDeltaWrite$PositionDeltaBatchWrite" + + /** True only for Iceberg's JVM DeltaWrite implementation. */ + def isIcebergPositionDeltaWrite(write: Any): Boolean = + write != null && tryLoadClass(positionDeltaWriteClassName).exists(_.isInstance(write)) + + /** True only for the BatchWrite enclosed by Iceberg's SparkPositionDeltaWrite. */ + def isIcebergPositionDeltaBatchWrite(batchWrite: Any): Boolean = + batchWrite != null && batchWrite.getClass.getName.startsWith( + positionDeltaBatchWriteClassPrefix) + def getOuterSparkWrite(batchWrite: Any): Option[Any] = { if (batchWrite == null) None else { @@ -1423,6 +1437,14 @@ object IcebergReflection extends Logging { def getTableFromSparkWrite(sparkWrite: Any): Option[Any] = getSparkWriteField(sparkWrite, "table") + /** Table owned by either Iceberg's regular SparkWrite or its position-delta write. */ + def getTableFromIcebergWrite(write: Any): Option[Any] = + if (isIcebergPositionDeltaWrite(write)) { + reflectField(write, "table") + } else { + getTableFromSparkWrite(write) + } + def getWritePropertiesFromSparkWrite(sparkWrite: Any): Option[Map[String, String]] = { import scala.jdk.CollectionConverters._ getSparkWriteField(sparkWrite, "writeProperties") @@ -2010,19 +2032,36 @@ object IcebergReflection extends Logging { } /** - * The locations of the data files carried by a `SparkWrite$TaskCommit` message (its - * package-private `files()`), or empty when `message` is not one. Used to clean up after a - * write job that failed before any commit was attempted. - */ - def taskCommitFileLocations(message: AnyRef): Seq[String] = + * The newly written file locations carried by an Iceberg task commit message. Plain + * `SparkWrite$TaskCommit` exposes `files()`; position-delta commits expose separate + * `dataFiles()` and `deleteFiles()` arrays. Do not include `rewrittenDeleteFiles()`: those + * files pre-date this job and are only candidates for removal after a successful delta commit. + * + * Used to clean completed tasks after a write job fails before any commit is attempted. Both + * Iceberg write implementations keep their own cleanup disabled in that failure mode. + */ + def taskCommitFileLocations(message: AnyRef): Seq[String] = { + def locations(methodName: String): Seq[String] = + findMethodInHierarchy(message.getClass, methodName).toSeq + .flatMap { method => + method.invoke(message) match { + case array: Array[_] => + array.toSeq.flatMap(file => extractFileLocation(file.getClass, file)) + case _ => Seq.empty + } + } + findMethodInHierarchy(message.getClass, "files") match { case Some(files) => files.invoke(message) match { - case array: Array[_] => array.toSeq.flatMap(f => extractFileLocation(f)) + case array: Array[_] => + array.toSeq.flatMap(file => extractFileLocation(file.getClass, file)) case _ => Seq.empty } - case None => Seq.empty + case None => + locations("dataFiles") ++ locations("deleteFiles") } + } /** The table's `FileIO` (`table.io()`). Iceberg requires `FileIO` to be `Serializable`. */ def getTableIO(table: Any): Option[AnyRef] = diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergStorageSchemes.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergStorageSchemes.scala new file mode 100644 index 00000000000..51f4b5f7019 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergStorageSchemes.scala @@ -0,0 +1,77 @@ +/* + * 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.iceberg + +import java.util.Locale + +import org.apache.spark.internal.Logging + +import org.apache.comet.NativeBase + +/** + * The storage schemes the native Iceberg storage factory publishes over JNI, so the JVM scan and + * write gates decline what native cannot open instead of failing at execution. Native is the only + * source of these lists. Every caller sits behind `isCometLoaded`, so a JVM where the library is + * not loaded never plans a native scan or write, and both sets are simply empty there. + */ +private[comet] object IcebergStorageSchemes extends Logging { + + lazy val read: Set[String] = load(forWrite = false) + lazy val write: Set[String] = load(forWrite = true) + + /** + * Splits the comma-joined native list, trimming and lowercasing each entry. A null or blank + * list throws: `builtin_storage_schemes` is a fixed non-empty constant natively, so a loaded + * library that publishes nothing can only be a build or marshalling bug, and a warning here + * would silently disable every native Iceberg scan and write. + */ + private[comet] def parse(joined: String, forWrite: Boolean): Set[String] = { + val schemes = Option(joined).toSeq + .flatMap(_.split(",")) + .map(_.trim.toLowerCase(Locale.ROOT)) + .filter(_.nonEmpty) + .toSet + if (schemes.isEmpty) { + val path = if (forWrite) "write" else "read" + throw new IllegalStateException( + s"Comet native library published no Iceberg $path scheme list") + } + schemes + } + + /** + * Loads one list over JNI, or answers the empty set with a warning when the library is not + * loaded; that branch never runs a plan, because every caller sits behind `isCometLoaded`. + * `isLoaded` is a parameter so the unloaded answer can be tested without unloading the library. + * Once loaded, anything but a non-empty list is a build bug that must fail loudly: a native + * fault while answering the probe propagates, and a blank answer throws in `parse`. + */ + private[comet] def load( + forWrite: Boolean, + isLoaded: Boolean = NativeBase.isLoaded): Set[String] = { + if (isLoaded) { + parse(NativeBase.icebergStorageSchemes(forWrite), forWrite) + } else { + val path = if (forWrite) "write" else "read" + logWarning(s"Comet native library is not loaded; the Iceberg $path scheme list is empty") + Set.empty + } + } +} diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteDispatch.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteDispatch.scala new file mode 100644 index 00000000000..970179753b4 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteDispatch.scala @@ -0,0 +1,28 @@ +/* + * 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.iceberg + +sealed trait IcebergWriteDispatch extends Product with Serializable + +case object PlainIcebergWrite extends IcebergWriteDispatch + +case class ReplaceDataWrite(info: ReplaceDataDispatchInfo) extends IcebergWriteDispatch + +case class PositionDeltaWrite(info: WriteDeltaDispatchInfo) extends IcebergWriteDispatch diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteLogical.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteLogical.scala index 6ceb526d4c9..c79e32365a1 100644 --- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteLogical.scala +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteLogical.scala @@ -30,7 +30,7 @@ case class IcebergWriteLogical( child: LogicalPlan, // Driver-side only: AQE re-planning is driver-local and write commands aren't cached. @transient batchWrite: BatchWrite, - replaceDataDispatch: Option[ReplaceDataDispatchInfo] = None) + dispatch: IcebergWriteDispatch = PlainIcebergWrite) extends UnaryNode { // Owns the commit-message attribute so the physical writer keeps the same exprId across diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala index 85fbcc8ce3a..dbb19fba797 100644 --- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala @@ -31,8 +31,8 @@ import org.apache.comet.CometSparkSessionExtensions.isCometLoaded import org.apache.comet.shims.ShimCometMergeRows /** - * Spark Strategy that intercepts Iceberg V2 copy-on-write logical writes and emits Comet's - * two-operator physical tree. + * Spark strategy for Comet's split Iceberg V2 writer. WriteDelta is modeled and dispatched here, + * but position-delta rows are still executed by Iceberg's JVM DeltaWriter. */ case class IcebergWriteStrategy(session: SparkSession) extends SparkStrategy { @@ -44,38 +44,68 @@ case class IcebergWriteStrategy(session: SparkSession) extends SparkStrategy { return Nil } // Planner strategies run before CometRule, so plan-only mode needs its own guard here. - if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(session.sessionState.conf)) { + if (CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(conf)) { return Nil } plan match { case ad: AppendData => - matchedSparkWrite(ad.table, ad.write, ad.query, replaceDataDispatch = None).toList + matchedSparkWrite(ad.table, ad.write, ad.query, PlainIcebergWrite).toList case obe: OverwriteByExpression => - matchedSparkWrite(obe.table, obe.write, obe.query, replaceDataDispatch = None).toList + matchedSparkWrite(obe.table, obe.write, obe.query, PlainIcebergWrite).toList case opd: OverwritePartitionsDynamic => - matchedSparkWrite(opd.table, opd.write, opd.query, replaceDataDispatch = None).toList + matchedSparkWrite(opd.table, opd.write, opd.query, PlainIcebergWrite).toList case rd: ReplaceData => matchedSparkWrite( rd.originalTable, rd.write, rd.query, - replaceDataDispatch = IcebergReplaceDataShim.extractProjections(rd)).toList - case plan if IcebergReflection.isReplaceIcebergData(plan) => + IcebergReplaceDataShim + .extractProjections(rd) + .map(ReplaceDataWrite) + .getOrElse(PlainIcebergWrite)).toList + case replace if IcebergReflection.isReplaceIcebergData(replace) => IcebergReflection - .extractReplaceIcebergDataFields(plan) + .extractReplaceIcebergDataFields(replace) .flatMap { case (_, query, originalTable, write) => matchedSparkWrite( originalTable.asInstanceOf[org.apache.spark.sql.catalyst.analysis.NamedRelation], write.asInstanceOf[Option[Write]], query.asInstanceOf[LogicalPlan], - replaceDataDispatch = None) + PlainIcebergWrite) } .toList // Hit by AQE. - case l @ IcebergWriteLogical(child, batchWrite, replaceDataDispatch) => - Seq(IcebergWriteExec(batchWrite, l.output, planLater(child), replaceDataDispatch)) - case _ => Nil + case l @ IcebergWriteLogical(child, batchWrite, dispatch) => + Seq(IcebergWriteExec(batchWrite, l.output, planLater(child), dispatch)) + case delta => + IcebergDeltaLogicalShim + .extract(delta) + .flatMap { fields => + fields.command.flatMap { command => + fields.write.flatMap { deltaWrite => + if (!IcebergReflection.isIcebergPositionDeltaWrite(deltaWrite)) { + None + } else { + WriteDeltaDispatchInfo + .build( + fields.projections, + fields.query.output, + IcebergDeltaWriterShim.OperationCodes, + Some(command)) + .flatMap { info => + buildDeltaTwoOp( + deltaWrite, + fields.originalTable, + fields.query, + PositionDeltaWrite(info), + Some(command)) + } + } + } + } + } + .toList } } @@ -83,12 +113,12 @@ case class IcebergWriteStrategy(session: SparkSession) extends SparkStrategy { table: org.apache.spark.sql.catalyst.analysis.NamedRelation, write: Option[Write], query: LogicalPlan, - replaceDataDispatch: Option[ReplaceDataDispatchInfo]): Option[SparkPlan] = { + dispatch: IcebergWriteDispatch): Option[SparkPlan] = { table match { case rel: DataSourceV2Relation => write.flatMap { w => if (IcebergReflection.isIcebergSparkWrite(w)) { - buildTwoOp(w, rel, query, replaceDataDispatch) + buildTwoOp(w, rel, query, dispatch) } else { None } @@ -104,32 +134,66 @@ case class IcebergWriteStrategy(session: SparkSession) extends SparkStrategy { * writer's child is wrapped in [[IcebergWriteLogical]] so AQE re-emits only the data-writing * operator on each re-plan as opposed to multiple new commit operators. * - * Iceberg's `SparkWrite` never asks for Spark's commit coordinator, so the - * `useCommitCoordinator` fallback below is defensive coverage in case a future Iceberg version - * changes that; the split writer's per-task commit protocol does not use it. + * Iceberg's current JVM writers do not ask for Spark's commit coordinator, so the + * `useCommitCoordinator` fallback is defensive coverage in case that changes; the split + * writer's per-task commit protocol does not use it. */ private def buildTwoOp( write: Write, rel: DataSourceV2Relation, query: LogicalPlan, - replaceDataDispatch: Option[ReplaceDataDispatchInfo]): Option[SparkPlan] = { + dispatch: IcebergWriteDispatch): Option[SparkPlan] = { val batchWrite = write.toBatch if (batchWrite.useCommitCoordinator()) { return None } - // To mirror Spark ReplaceData semantics we invalidate our cache of the state of - // `originalTable`. + // To mirror Spark's V2 write semantics, invalidate the cached state of `originalTable`. val refresh: () => Unit = () => IcebergRefreshCacheShim.refreshCache(session, rel) Some( - IcebergCommitPlanShim.wrap(IcebergCommitExec( - batchWrite, - write, - refresh, - // `replaceDataDispatch` may project the data into the format the writer expects. - planLater( - IcebergWriteLogical( - ShimCometMergeRows.withNativeMergeSummary(query), + IcebergCommitPlanShim.wrap( + IcebergCommitExec( + batchWrite, + write, + refresh, + planLater( + IcebergWriteLogical( + ShimCometMergeRows.withNativeMergeSummary(query), + batchWrite, + dispatch))))) + } + + /** + * Builds the same split commit/write tree for Iceberg's position-delta writer. The task-side + * writer remains Iceberg's JVM `DeltaWriter`; this only lets Comet own the split-plan + * lifecycle, AQE re-planning, reporting, and driver commit. + */ + private def buildDeltaTwoOp( + write: Write, + table: org.apache.spark.sql.catalyst.analysis.NamedRelation, + query: LogicalPlan, + dispatch: IcebergWriteDispatch, + command: Option[DeltaCommand]): Option[SparkPlan] = { + table match { + case rel: DataSourceV2Relation => + val batchWrite = write.toBatch + if (!IcebergReflection.isIcebergPositionDeltaBatchWrite(batchWrite) || + batchWrite.useCommitCoordinator()) { + None + } else { + val refresh: () => Unit = () => IcebergRefreshCacheShim.refreshCache(session, rel) + val commit = IcebergCommitExec( batchWrite, - replaceDataDispatch))))) + write, + refresh, + planLater( + IcebergWriteLogical( + ShimCometMergeRows.withNativeMergeSummary(query), + batchWrite, + dispatch)), + command) + Some(IcebergCommitPlanShim.wrap(commit)) + } + case _ => None + } } } diff --git a/spark/src/main/scala/org/apache/comet/iceberg/WriteDeltaDispatchInfo.scala b/spark/src/main/scala/org/apache/comet/iceberg/WriteDeltaDispatchInfo.scala new file mode 100644 index 00000000000..4aaf43aebec --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/iceberg/WriteDeltaDispatchInfo.scala @@ -0,0 +1,64 @@ +/* + * 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.iceberg + +import org.apache.spark.sql.catalyst.ProjectingInternalRow +import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.catalyst.util.{RowDeltaUtils, WriteDeltaProjections} +import org.apache.spark.sql.types.IntegerType + +case class DeltaOperationCodes(delete: Int, update: Int, insert: Int, reinsert: Option[Int]) + +/** + * Version-neutral task contract for Spark's operation-coded WriteDelta row stream. + * + * This intentionally contains only state consumed by Iceberg's JVM DeltaWriter path. Native + * position-delete schemas and metadata layouts belong to the native WriteDelta implementation. + */ +case class WriteDeltaDispatchInfo( + operationOrdinal: Int, + operationCodes: DeltaOperationCodes, + rowProjection: Option[ProjectingInternalRow], + rowIdProjection: ProjectingInternalRow, + metadataProjection: Option[ProjectingInternalRow], + command: Option[DeltaCommand] = None) + +private[iceberg] object WriteDeltaDispatchInfo { + + def build( + projections: WriteDeltaProjections, + childOutput: Seq[Attribute], + operationCodes: DeltaOperationCodes, + command: Option[DeltaCommand] = None): Option[WriteDeltaDispatchInfo] = { + childOutput.headOption.collect { + case attribute + if attribute.name == RowDeltaUtils.OPERATION_COLUMN && + attribute.dataType == IntegerType && + !attribute.nullable => + WriteDeltaDispatchInfo( + operationOrdinal = 0, + operationCodes = operationCodes, + rowProjection = projections.rowProjection, + rowIdProjection = projections.rowIdProjection, + metadataProjection = projections.metadataProjection, + command = command) + } + } +} 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 c16d1fc8f8f..a47f1c2a6b4 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -25,7 +25,7 @@ import scala.collection.mutable.ListBuffer import scala.jdk.CollectionConverters._ import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.expressions.{Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, FloatLiteral, GreaterThan, GreaterThanOrEqual, KnownFloatingPointNormalized, LessThan, LessThanOrEqual, NamedExpression, Remainder, SortOrder} +import org.apache.spark.sql.catalyst.expressions.{Divide, DoubleLiteral, Expression, FloatLiteral, KnownFloatingPointNormalized, NamedExpression, Remainder, SortOrder} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, Final, Partial, PartialMerge} import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero import org.apache.spark.sql.catalyst.rules.Rule @@ -186,6 +186,11 @@ object CometExecRule { */ private val TYPED_DATASET_PARTIAL_READER: TreeNodeTag[String] = TreeNodeTag[String]("comet.typedDatasetPartialReader") + + /** Why Comet does not convert rows to Arrow in a plan that reads `InputFileBlockHolder`. */ + private val INPUT_FILE_BLOCK_FALLBACK_REASON: String = + "Spark to Arrow conversion is not compatible with input_file_name, " + + "input_file_block_start, or input_file_block_length" } /** @@ -205,11 +210,14 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) /** * Revert any `CometShuffleExchangeExec` with `CometColumnarShuffle` whose parent and child are - * both non-Comet `HashAggregateExec` / `ObjectHashAggregateExec` operators back to the original - * Spark `ShuffleExchangeExec`. This is the partial-final-aggregate pattern where Comet couldn't + * both non-Comet aggregate operators (any `BaseAggregateExec`) back to the original Spark + * `ShuffleExchangeExec`. This is the partial-final-aggregate pattern where Comet couldn't * convert either aggregate; keeping a columnar shuffle between them only adds * row->arrow->shuffle->arrow->row conversion overhead with no Comet consumer on either side. - * See https://github.com/apache/datafusion-comet/issues/4004. + * See https://github.com/apache/datafusion-comet/issues/4004. A native shuffle over rows that + * `convertShuffleInput` converted is the same pattern, and is reverted the same way. A + * `SortAggregateExec` with grouping keys reads the shuffle through the `SortExec` that + * `EnsureRequirements` adds, so it is not the shuffle's parent and the shuffle stays. * * The match is intentionally narrow (both sides must be row-based aggregates that remained JVM * after the main transform pass). Running the revert post-transform means we only fire when the @@ -230,26 +238,35 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) private def revertRedundantColumnarShuffle(plan: SparkPlan): SparkPlan = { def isAggregate(p: SparkPlan): Boolean = p.isInstanceOf[BaseAggregateExec] - def isRedundantShuffle(child: SparkPlan): Boolean = child match { - case s: CometShuffleExchangeExec => - s.shuffleType == CometColumnarShuffle && isAggregate(s.child) - case _ => false + // The Spark operator whose rows a Comet shuffle writes: the child of a JVM columnar shuffle, + // or the input of the conversion that `convertShuffleInput` put under a native shuffle. + def rowInput(s: CometShuffleExchangeExec): Option[SparkPlan] = + (s.shuffleType, s.child) match { + case (CometColumnarShuffle, child) => Some(child) + case (CometNativeShuffle, conversion: CometSparkToColumnarExec) => Some(conversion.child) + case _ => None + } + + def redundantShuffleInput(child: SparkPlan): Option[SparkPlan] = child match { + case s: CometShuffleExchangeExec => rowInput(s).filter(isAggregate) + case _ => None } plan.transform { - case op if isAggregate(op) && op.children.exists(isRedundantShuffle) => - val newChildren = op.children.map { - case s: CometShuffleExchangeExec - if s.shuffleType == CometColumnarShuffle && isAggregate(s.child) => - val reverted = - s.originalPlan.withNewChildren(Seq(s.child)).asInstanceOf[ShuffleExchangeExec] - reverted.setTagValue(CometExecRule.SKIP_COMET_SHUFFLE_TAG, ()) - logInfo( - "Reverting Comet columnar shuffle to Spark shuffle between " + - s"${op.getClass.getSimpleName} and ${s.child.getClass.getSimpleName} " + - "(no Comet operator on either side to consume columnar output)") - reverted - case other => other + case op if isAggregate(op) && op.children.exists(redundantShuffleInput(_).isDefined) => + val newChildren = op.children.map { child => + (child, redundantShuffleInput(child)) match { + case (s: CometShuffleExchangeExec, Some(input)) => + val reverted = + s.originalPlan.withNewChildren(Seq(input)).asInstanceOf[ShuffleExchangeExec] + reverted.setTagValue(CometExecRule.SKIP_COMET_SHUFFLE_TAG, ()) + logInfo( + "Reverting Comet shuffle to Spark shuffle between " + + s"${op.getClass.getSimpleName} and ${input.getClass.getSimpleName} " + + "(no Comet operator on either side to consume columnar output)") + reverted + case _ => child + } } op.withNewChildren(newChildren) } @@ -397,6 +414,9 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) */ // spotless:on private def transform(plan: SparkPlan): SparkPlan = { + // Walks the whole plan, so it is lazy: only consulted once a node could be converted. + lazy val readsInputFileBlock = CometScanRule.readsInputFileBlock(plan) + if (CometConf.COMET_CONVERT_FROM_TYPED_DATASET_ENABLED.get(conf)) { tagPartiallyReadTypedDatasetOutputs(plan) } @@ -472,7 +492,7 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) s"Set ${CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key}=true to enable it.") } - if (shouldApplySparkToColumnar(conf, scan)) { + if (shouldApplySparkToColumnar(conf, scan, readsInputFileBlock)) { convertToComet(scan, CometSparkToColumnarExec).getOrElse(scan) } else { scan @@ -494,7 +514,7 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) case op: LeafExecNode if hasEnabledHandler(op) => convertToComet(op, allExecs(op.getClass)) .orElse { - if (shouldApplySparkToColumnar(conf, op)) { + if (shouldApplySparkToColumnar(conf, op, readsInputFileBlock)) { convertToComet(op, CometSparkToColumnarExec) } else { None @@ -502,7 +522,7 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) } .getOrElse(op) - case op if shouldApplySparkToColumnar(conf, op) => + case op if shouldApplySparkToColumnar(conf, op, readsInputFileBlock) => convertToComet(op, CometSparkToColumnarExec).getOrElse(op) // Typed Dataset operations (`map`, `flatMap`, `mapPartitions`, `mapGroups`, ...) pass JVM @@ -518,6 +538,10 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) "Comet does not convert the output of a typed Dataset operation when " + s"$reader can stop reading it early, because filling an Arrow batch would " + "run the user function on rows that Spark never reaches") + // The conversion reads ahead of input_file_name and friends, as it does over a leaf. + // See shouldApplySparkToColumnar. + case None if readsInputFileBlock => + withFallbackReason(op, CometExecRule.INPUT_FILE_BLOCK_FALLBACK_REASON) case None => convertTypedDatasetOutput(op) } @@ -542,23 +566,15 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) // DataWritingCommandExec and re-implement the write framework inside CometNativeWriteExec. // This path is retained only for 3.4/3.5 and goes away with them. // - // AQE reoptimization looks for `DataWritingCommandExec` or `WriteFilesExec` - // if there is none it would reinsert write nodes, and since Comet remap those nodes - // to Comet counterparties the write nodes are twice to the plan. - // Checking if AQE inserted another write Command on top of existing write command - case _ @DataWritingCommandExec(_, w: WriteFilesExec) - if !isSpark40Plus && w.child.isInstanceOf[CometNativeWriteExec] => - w.child - + // `originalPlan` is that command. This rule copies `originalPlan.logicalLink` onto the + // Comet node, so AQE re-plans the write with the command rather than with whatever child + // happened to sit under it. No second DataWritingCommandExec is inserted on top. case op: DataWritingCommandExec if !isSpark40Plus => convertToComet(op, CometDataWritingCommand).getOrElse(op) - // AQE re-fires the Iceberg write planning on every stage materialisation, so a - // partitioned write's physical sub-tree may already contain a `CometIcebergWriteExec`. - // Unwrap to avoid a double conversion. - case op: IcebergWriteExec if op.child.isInstanceOf[CometIcebergWriteExec] => - op.child - + // `originalPlan` is this IcebergWriteExec, so AQE re-plans the write as this node. + // A shuffle directly under the native write stays in the child stage and is not wrapped + // again. case op: IcebergWriteExec if CometConf.COMET_ICEBERG_NATIVE_WRITE_ENABLED.get(op.conf) => convertToComet(op, CometIcebergNativeWrite).getOrElse(op) @@ -616,6 +632,9 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) case s: ShuffleExchangeExec if shouldSkipCometShuffle(s) => preserveSparkAggregateBuffers(s) + case s: ShuffleExchangeExec if CometShuffleExchangeExec.convertsInputForNativeShuffle(s) => + convertShuffleInput(s) + case s: ShuffleExchangeExec => convertToComet(s, CometShuffleExchangeExec) .getOrElse(preserveSparkAggregateBuffers(s)) @@ -643,7 +662,8 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) op case _: ColumnarToRowTransition => // A transition does no work of its own. This rule only meets one that - // `convertTypedDatasetOutput` inserted on an earlier pass over the same plan. + // `convertTypedDatasetOutput` or `convertShuffleInput` inserted on an earlier pass + // over the same plan. op case _: WriteFilesExec => // The write is converted at the enclosing DataWritingCommandExec above: on Spark 3.x @@ -887,58 +907,6 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) } } - private def normalizePlan(plan: SparkPlan): SparkPlan = { - plan.transformUp { - case p: ProjectExec => - val newProjectList = p.projectList.map(normalize(_).asInstanceOf[NamedExpression]) - ProjectExec(newProjectList, p.child) - case f: FilterExec => - val newCondition = normalize(f.condition) - FilterExec(newCondition, f.child) - } - } - - // Spark will normalize NaN and zero for floating point numbers for several cases. - // See `NormalizeFloatingNumbers` optimization rule in Spark. - // However, one exception is for comparison operators. Spark does not normalize NaN and zero - // because they are handled well in Spark (e.g., `SQLOrderingUtil.compareFloats`). But the - // comparison functions in arrow-rs do not normalize NaN and zero. So we need to normalize NaN - // and zero for comparison operators in Comet. - private def normalize(expr: Expression): Expression = { - expr.transformUp { - case EqualTo(left, right) => - EqualTo(normalizeNaNAndZero(left), normalizeNaNAndZero(right)) - case EqualNullSafe(left, right) => - EqualNullSafe(normalizeNaNAndZero(left), normalizeNaNAndZero(right)) - case GreaterThan(left, right) => - GreaterThan(normalizeNaNAndZero(left), normalizeNaNAndZero(right)) - case GreaterThanOrEqual(left, right) => - GreaterThanOrEqual(normalizeNaNAndZero(left), normalizeNaNAndZero(right)) - case LessThan(left, right) => - LessThan(normalizeNaNAndZero(left), normalizeNaNAndZero(right)) - case LessThanOrEqual(left, right) => - LessThanOrEqual(normalizeNaNAndZero(left), normalizeNaNAndZero(right)) - case Divide(left, right, evalMode) => - Divide(left, normalizeNaNAndZero(right), evalMode) - case Remainder(left, right, evalMode) => - Remainder(left, normalizeNaNAndZero(right), evalMode) - } - } - - private def normalizeNaNAndZero(expr: Expression): Expression = { - expr match { - case _: KnownFloatingPointNormalized => expr - case FloatLiteral(f) if !f.isNaN && !f.equals(-0.0f) => expr - case DoubleLiteral(d) if !d.isNaN && !d.equals(-0.0d) => expr - case _ => - expr.dataType match { - case _: FloatType | _: DoubleType => - KnownFloatingPointNormalized(NormalizeNaNAndZero(expr)) - case _ => expr - } - } - } - /** * A relation keeps the cache format it was stored in, since `spark.sql.cache.serializer` is * static, so a plan that runs without Comet's native execution still reads relations cached in @@ -969,6 +937,41 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) newPlan } + /** + * Wraps the divisor of a floating-point `Divide` or `Remainder` in a Project or Filter in + * `NormalizeNaNAndZero`. Native comparisons already follow Spark's ordering, including the + * check for a zero divisor, so the quotient does not need this. It also canonicalizes a NaN + * divisor, though, and with it the NaN quotient: `1.0D / (-d)` of a NaN `d` is otherwise a NaN + * with the sign bit set, which `percentile_approx` still orders below every other value + * (https://github.com/apache/datafusion-comet/issues/6519). Remove this once that is fixed. + */ + private def normalizeDivisors(plan: SparkPlan): SparkPlan = { + def normalize(expr: Expression): Expression = expr.transformUp { + case Divide(left, right, evalMode) => Divide(left, normalizeNaNAndZero(right), evalMode) + case Remainder(left, right, evalMode) => + Remainder(left, normalizeNaNAndZero(right), evalMode) + } + plan.transformUp { + case p: ProjectExec => + ProjectExec(p.projectList.map(normalize(_).asInstanceOf[NamedExpression]), p.child) + case f: FilterExec => FilterExec(normalize(f.condition), f.child) + } + } + + private def normalizeNaNAndZero(expr: Expression): Expression = { + expr match { + case _: KnownFloatingPointNormalized => expr + case FloatLiteral(f) if !f.isNaN && !f.equals(-0.0f) => expr + case DoubleLiteral(d) if !d.isNaN && !d.equals(-0.0d) => expr + case _ => + expr.dataType match { + case _: FloatType | _: DoubleType => + KnownFloatingPointNormalized(NormalizeNaNAndZero(expr)) + case _ => expr + } + } + } + private def _apply(plan: SparkPlan): SparkPlan = { // We shouldn't transform Spark query plan if Comet is not loaded. if (!isCometLoaded(conf)) { @@ -989,12 +992,15 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) plan } } else { - val normalizedPlan = normalizePlan(plan) + val normalizedPlan = normalizeDivisors(plan) + // The rewrite removes the sorts of each join it converts, so put back any sort a kept + // operator above still needs. This runs before transform(), which converts the sorts it + // adds like any other SortExec. val planWithJoinRewritten = if (CometConf.COMET_FORCE_SHJ.get()) { - normalizedPlan.transformUp { case p => - RewriteJoin.rewrite(p) - } + RewriteJoin.restoreRequiredOrdering(normalizedPlan.transformUp { case p => + RewriteJoin.rewrite(p, conf) + }) } else { normalizedPlan } @@ -1405,7 +1411,51 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) } } - private def shouldApplySparkToColumnar(conf: SQLConf, op: SparkPlan): Boolean = { + /** + * Converts the rows a Spark operator feeds to a shuffle to Arrow, so that the shuffle runs as + * native shuffle instead of the JVM columnar shuffle. See + * [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]], and + * [[CometShuffleExchangeExec.convertsInputForNativeShuffle]] for when it applies. + * + * Spark inserts the columnar transitions after this rule, but it does not look below a + * `RowToColumnarTransition` such as `CometSparkToColumnarExec`. Without a transition, a Spark + * operator in the child's subtree that reads a Comet operator would do so through + * `CometExec.doExecute`, Spark's interpreted columnar-to-row path. So the subtree gets its + * transitions now, from Spark's own rule, and `EliminateRedundantTransitions` later replaces + * each one over a Comet child with Comet's own. Spark's rule leaves existing transitions alone, + * which matters because this rule runs over the same plan twice under AQE. + */ + private def convertShuffleInput(s: ShuffleExchangeExec): SparkPlan = { + val child = + ApplyColumnarRulesAndInsertTransitions(Seq.empty, outputsColumnar = false).apply(s.child) + convertToComet(child, CometSparkToColumnarExec) + .flatMap(converted => + convertToComet(s.withNewChildren(Seq(converted)), CometShuffleExchangeExec)) + .getOrElse( + convertToComet(s, CometShuffleExchangeExec).getOrElse(preserveSparkAggregateBuffers(s))) + } + + private def shouldApplySparkToColumnar( + conf: SQLConf, + op: SparkPlan, + readsInputFileBlock: => Boolean): Boolean = { + // A converted leaf reads ahead of the Spark operator that evaluates input_file_name and + // friends: the conversion fills a whole batch, and Comet operators above it may pull more + // before they emit. By the time Spark evaluates them, the leaf's reader may have moved on to + // a later file or unset InputFileBlockHolder at the end of its input, so rows would report + // another file's values or the unset defaults. Leave the leaf on Spark so that the plan above + // it stays on Spark too. + if (!canApplySparkToColumnar(conf, op)) { + false + } else if (readsInputFileBlock) { + withFallbackReason(op, CometExecRule.INPUT_FILE_BLOCK_FALLBACK_REASON) + false + } else { + true + } + } + + private def canApplySparkToColumnar(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 // columnar batches, such as Spark's vectorized readers, will also be converted to native diff --git a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala index b109fa158ed..c7c9a31909d 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala @@ -47,7 +47,7 @@ import org.apache.spark.sql.types._ import org.apache.comet.{CometConf, DataTypeSupport, NativeBase} import org.apache.comet.CometConf._ import org.apache.comet.CometSparkSessionExtensions.{isCometLoaded, isSpark35Plus, withFallbackReason, withFallbackReasons} -import org.apache.comet.iceberg.{CometIcebergNativeScanMetadata, IcebergReflection} +import org.apache.comet.iceberg.{CometIcebergNativeScanMetadata, IcebergReflection, IcebergStorageSchemes} import org.apache.comet.objectstore.NativeConfig import org.apache.comet.parquet.CometParquetUtils.{encryptionEnabled, isEncryptionConfigSupported, readFieldId} import org.apache.comet.serde.operator.{CometIcebergNativeScan, CometNativeScan} @@ -354,14 +354,10 @@ case class CometScanRule(session: SparkSession) withFallbackReason(scanExec, "Native Parquet Variant scans do not support encryption") return None } - // input_file_name, input_file_block_start, and input_file_block_length read from - // InputFileBlockHolder, a thread-local set by Spark's FileScanRDD. The native DataFusion - // scan does not use FileScanRDD, so these expressions would return empty/default values. - if (plan.exists(node => - node.expressions.exists(_.exists { - case _: InputFileName | _: InputFileBlockStart | _: InputFileBlockLength => true - case _ => false - }))) { + // input_file_name, input_file_block_start, and input_file_block_length read values that + // Spark's FileScanRDD sets. The native DataFusion scan does not use FileScanRDD, so these + // expressions would return empty/default values. + if (CometScanRule.readsInputFileBlock(plan)) { withFallbackReason( scanExec, "Native Parquet scan is not compatible with input_file_name, " + @@ -1185,6 +1181,20 @@ case class CometScanTypeChecker() extends DataTypeSupport with CometTypeShim { object CometScanRule extends Logging { + /** + * Whether any node in `plan` evaluates `input_file_name`, `input_file_block_start` or + * `input_file_block_length`. These read `InputFileBlockHolder`, a thread-local that the reader + * producing the rows (`FileScanRDD`, the V2 file readers, `HadoopRDD`, `NewHadoopRDD`, + * connector readers such as Iceberg's) sets as it moves from file to file, so they only return + * each row's values when Spark evaluates them as that reader produces the row. + */ + def readsInputFileBlock(plan: SparkPlan): Boolean = + plan.exists(node => + node.expressions.exists(_.exists { + case _: InputFileName | _: InputFileBlockStart | _: InputFileBlockLength => true + case _ => false + })) + // Memo of `NativeBase.isObjectStoreSchemeSupported`, keyed by the probe URL rather than the // scheme: object_store's parser keys on (scheme, host-presence), so an authorityless URL would // otherwise poison the authority-bearing form of the same scheme. @@ -1276,20 +1286,25 @@ object CometScanRule extends Logging { org.apache.spark.sql.catalyst.trees.TreeNodeTag[Unit]("comet.skipCometScan") /** - * Schemes Comet's native Iceberg scan can actually open, mirroring the match arms in - * `native/core/src/execution/operators/iceberg_common.rs::storage_factory_for`. Deliberately - * NOT delegated to `isNativelyReadableScheme`: object_store recognizes schemes (http/https, - * azure, memory) that iceberg-rust's OpenDAL storage factory cannot build, and admitting them - * here turns a clean JVM fallback into a native runtime "Unsupported storage scheme" error. Add - * here what you add to `storage_factory_for` (currently Aliyun `oss` and GCS `gs`). - * S3-compliant aliases like `blob` are opt-in via `fs.comet.s3Compliant.schemes` (see - * `isIcebergReadableScheme`), not hardcoded, since the native planner opens them via S3. The - * write path keeps its own list (`CometIcebergNativeWrite.SupportedStorageSchemes`), which - * differs deliberately: it excludes `oss` (fails closed, see `storage_factory_for`) and - * includes `memory`. + * Schemes Comet's native Iceberg scan can open, loaded from the native storage factory over JNI + * so this gate cannot drift from `storage_factory_for`. Lazy so that constructing the rule does + * not touch the native library before `isCometLoaded` has been consulted. Opt-in aliases from + * `fs.comet.s3Compliant.schemes` are additive (see `isIcebergReadableScheme`); the write path + * loads its own set (`CometIcebergNativeWrite.SupportedStorageSchemes`). + */ + private lazy val icebergReadableSchemes: Set[String] = IcebergStorageSchemes.read + + /** + * True when the Iceberg scan gate admits `scheme`, written exactly as recorded. Native opens a + * location by its raw scheme, and OpenDAL's S3 backend checks it against a `scheme://bucket/` + * prefix whose scheme comes from `Url::parse` and so is lowercase: `S3://` is not `s3://`, and + * `BLOB://` is not an opted-in `blob`. The alias set is lowercase + * (`NativeConfig.parseSchemeSet`), so an alias is admitted only when the location writes it in + * lowercase. The Parquet gate stays case-insensitive because it rewrites alias URLs to `s3://` + * before anything opens them. */ - private val icebergReadableSchemes: Set[String] = - Set("file", "s3", "s3a", "gs", "oss") + private def isAdmittedIcebergScheme(scheme: String, s3CompliantSchemes: Set[String]): Boolean = + icebergReadableSchemes.contains(scheme) || s3CompliantSchemes.contains(scheme) /** * "Supported schemes: ..." suffix shared by the Iceberg scheme-fallback messages. Lists the @@ -1310,8 +1325,7 @@ object CometScanRule extends Logging { s3CompliantSchemes: Set[String]): Boolean = { val scheme = uri.getScheme if (scheme == null) return true - val lower = scheme.toLowerCase(Locale.ROOT) - icebergReadableSchemes.contains(lower) || s3CompliantSchemes.contains(lower) + isAdmittedIcebergScheme(scheme, s3CompliantSchemes) } /** @@ -1373,9 +1387,6 @@ object CometScanRule extends Logging { // hasOpenableAuthority); non-empty => decline. One example suffices for the message. var hostlessLocation: Option[String] = None - // Union of the two admitted scheme sets, built once so the per-file loop does one lookup. - val openableSchemes = icebergReadableSchemes ++ s3CompliantSchemes - // Classify one data/delete file location; see `icebergReadableSchemes` for why that allowlist // is narrower than the Parquet native gate. Runs per data and delete file, so the scheme and // the bucket are each derived once and threaded down. @@ -1386,9 +1397,8 @@ object CometScanRule extends Logging { // A schemeless local path routes to iceberg-rust's LocalFs and needs no host. val scheme = uri.getScheme if (scheme == null) return - val lower = scheme.toLowerCase(Locale.ROOT) - if (!openableSchemes.contains(lower)) { - unsupportedSchemes += lower + if (!isAdmittedIcebergScheme(scheme, s3CompliantSchemes)) { + unsupportedSchemes += scheme } else if (!hasOpenableAuthority(uri, s3CompliantSchemes)) { if (hostlessLocation.isEmpty) hostlessLocation = Some(rawPath) } else { diff --git a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala index 1181a573bc3..9ad2d003841 100644 --- a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala +++ b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala @@ -23,7 +23,8 @@ import org.apache.spark.internal.Logging import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial, PartialMerge} import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.comet.{CometBaseAggregateExec, CometColumnarToRowExec, CometExec, CometIcebergNativeScanExec, CometLocalTopKExec, CometNativeColumnarToRowExec, CometNativeScanExec, CometSparkToColumnarExec} +import org.apache.spark.sql.comet.{CometBaseAggregateExec, CometColumnarToRowExec, CometExec, CometNativeColumnarToRowExec, CometSparkToColumnarExec} +import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.execution.{ColumnarToRowExec, ColumnarToRowTransition, RowToColumnarExec, SparkPlan} import org.apache.spark.sql.execution.adaptive.QueryStageExec import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ShuffleExchangeLike} @@ -62,32 +63,66 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan plan match { case _: BroadcastExchangeLike => plan case exchange: ShuffleExchangeLike => - revertStageIfNeeded(exchange.child, exchange.supportsColumnar) - .map(reverted => exchange.withNewChildren(Seq(reverted))) - .getOrElse(plan) + revertShuffleStageIfNeeded(exchange).getOrElse(plan) case _ => - // Result stage: its output is collected as rows, so no consumer requires columnar input - // and the reverted stage needs no trailing R2C. + // Result stage: its output is collected as rows. revertStageIfNeeded(plan, outputColumnar = false).getOrElse(plan) } } private def applyForNonAQE(plan: SparkPlan): SparkPlan = { val withRevertedStages = plan.transformUp { case exchange: ShuffleExchangeLike => - revertStageIfNeeded(exchange.child, exchange.supportsColumnar) - .map(reverted => exchange.withNewChildren(Seq(reverted))) - .getOrElse(exchange) + revertShuffleStageIfNeeded(exchange).getOrElse(exchange) } revertStageIfNeeded(withRevertedStages, outputColumnar = false) .getOrElse(withRevertedStages) } /** - * Reverts the stage if C2R count exceeds threshold. Wraps in R2C if exchange needs columnar. + * Reverts the stage below a shuffle if needed and returns the shuffle over the reverted stage. + * A native shuffle over rows that `spark.comet.convert.shuffleInput.enabled` converted loses + * the conversion with the rest of the stage, so it goes back to the JVM columnar shuffle that + * the conversion replaced, which reads the stage's rows. Any other native shuffle consumes + * Arrow-backed Comet vectors, so the reverted stage is bridged back to them. + */ + private def revertShuffleStageIfNeeded(exchange: ShuffleExchangeLike): Option[SparkPlan] = + exchange match { + case s: CometShuffleExchangeExec + if s.shuffleType == CometNativeShuffle && convertsSparkRows(s.child) => + revertStageIfNeeded(s.child, s.supportsColumnar).map { reverted => + val columnar = s.copy(child = reverted, shuffleType = CometColumnarShuffle) + columnar.copyTagsFrom(s) + columnar + } + case _ => + val outputArrow = exchange match { + case comet: CometShuffleExchangeExec => comet.shuffleType == CometNativeShuffle + case _ => false + } + revertStageIfNeeded(exchange.child, exchange.supportsColumnar, outputArrow) + .map(reverted => exchange.withNewChildren(Seq(reverted))) + } + + /** + * Whether `plan` is the conversion that `spark.comet.convert.shuffleInput.enabled` puts over a + * Spark operator. A Comet transition directly under a `CometSparkToColumnarExec` is a stacked + * bridge that the revert strips and then restores for a native shuffle, not that conversion. + */ + private def convertsSparkRows(plan: SparkPlan): Boolean = plan match { + case conversion: CometSparkToColumnarExec => + !conversion.child.isInstanceOf[CometNativeColumnarToRowExec] && + !conversion.child.isInstanceOf[CometColumnarToRowExec] + case _ => false + } + + /** + * Reverts the stage if C2R count exceeds threshold, restoring the stage's output format when + * the reverted root does not satisfy it. */ private def revertStageIfNeeded( stagePlan: SparkPlan, - outputColumnar: Boolean): Option[SparkPlan] = { + outputColumnar: Boolean, + outputArrow: Boolean = false): Option[SparkPlan] = { val transitionCount = countTransitions(stagePlan) if (transitionCount <= maxTransitions) return None @@ -101,11 +136,27 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan val reason = s"Stage reverted: $transitionCount C2R transitions exceed threshold $maxTransitions" - val reverted = revertToSpark(stagePlan) - val result = if (outputColumnar && !reverted.supportsColumnar) { - RowToColumnarExec(withFallbackReason(reverted, reason)) + val reverted = + try { + revertToSpark(stagePlan) + } catch { + case e: CometExec.InvalidSparkFallbackException => + logWarning( + "Skipping transition-heavy stage reversion because a Comet operator could not " + + s"restore its Spark plan: ${e.getMessage}") + return None + } + val revertedWithReason = withFallbackReason(reverted, reason) + val result = if (outputArrow) { + // Native shuffle consumes Arrow-backed Comet vectors, not arbitrary Spark columnar + // batches. This bridge converts both row-based and vectorized Spark fallback roots. + CometSparkToColumnarExec(revertedWithReason) + } else if (outputColumnar && !reverted.supportsColumnar) { + RowToColumnarExec(revertedWithReason) + } else if (!outputColumnar && reverted.supportsColumnar) { + ColumnarToRowExec(revertedWithReason) } else { - withFallbackReason(reverted, reason) + revertedWithReason } Some(result) } @@ -146,16 +197,31 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan } /** - * Like `transformDown`, never descends stage-boundary children. + * Like `transformDown`, never descends stage-boundary children. If the rule rewrites the + * current node, re-apply it to the result so stacked transitions such as + * `CometSparkToColumnarExec(CometNativeColumnarToRowExec(x))` are fully unwrapped before + * children are visited. Spark's `transformDown` does not do this; leaving the inner C2R in + * place later calls `CometNativeColumnarToRowExec.withNewChildren` with a reverted row-based + * child, which asserts `child.supportsColumnar`. + * + * A rewrite can itself be the stage boundary. Unwrapping a transition that sits directly on a + * shuffle yields that shuffle, and descending into it strips transitions in the next stage. + * `transformStageUp` and `insertTransitions` do not cross the exchange, so those transitions + * would not be restored (#6152). Return the boundary unchanged. */ private def transformStageDown(plan: SparkPlan)( rule: PartialFunction[SparkPlan, SparkPlan]): SparkPlan = { val transformed = rule.applyOrElse(plan, identity[SparkPlan]) - val newChildren = transformed.children.map { child => - if (isStageBoundary(child)) child else transformStageDown(child)(rule) + if (transformed ne plan) { + if (isStageBoundary(transformed)) transformed + else transformStageDown(transformed)(rule) + } else { + val newChildren = transformed.children.map { child => + if (isStageBoundary(child)) child else transformStageDown(child)(rule) + } + if (newChildren == transformed.children) transformed + else transformed.withNewChildren(newChildren) } - if (newChildren == transformed.children) transformed - else transformed.withNewChildren(newChildren) } /** Like `transformUp`, never descends stage-boundary children. */ @@ -184,7 +250,27 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan count } + /** + * Checks for Comet operators whose original Spark plan is also their input. This must run + * before any bottom-up rewrite replaces children. Otherwise a stable `originalPlan` reference + * can keep pointing at the old Comet child after `withNewChildren`, hiding the alias and + * causing fallback to reconstruct that Comet child instead of a Spark operator. + */ + private def validateOriginalPlanAliases(plan: SparkPlan): Unit = plan match { + case _ if isStageBoundary(plan) => () + case cometExec: CometExec => + val sparkPlan = cometExec.originalPlan + if (sparkPlan != null && cometExec.children.exists(_ eq sparkPlan)) { + throw new CometExec.InvalidSparkFallbackException( + s"${cometExec.getClass.getSimpleName} aliases its original Spark plan with a child") + } + cometExec.children.foreach(validateOriginalPlanAliases) + case _ => + plan.children.foreach(validateOriginalPlanAliases) + } + private[rules] def revertToSpark(plan: SparkPlan): SparkPlan = { + validateOriginalPlanAliases(plan) val stripped = transformStageDown(plan) { case CometNativeColumnarToRowExec(child) => child case CometColumnarToRowExec(child) => child @@ -192,40 +278,12 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan case sparkToColumnar: CometSparkToColumnarExec => sparkToColumnar.child case RowToColumnarExec(child) => child } - val reverted = transformStageUp(stripped) { - // Local candidate selection was inserted by Comet. Only the outer TopK owns - // the original Spark operator's offset and projection. - case local: CometLocalTopKExec => local.child - case cometExec: CometExec => - if (cometExec.originalPlan.children.size == cometExec.children.size) { - val originalWithCurrentExpressions = cometExec match { - case scan: CometNativeScanExec => - // AQE's query-stage optimizer rewrites DPP placeholders in the live Comet scan - // before this post-columnar rule runs. The frozen FileSourceScanExec in - // originalPlan still contains SubqueryAdaptiveBroadcastExec, which cannot execute. - // Preserve the rewritten filters when reverting the scan to Spark. - val originalScan = scan.originalPlan.copy( - partitionFilters = scan.partitionFilters, - dataFilters = scan.dataFilters) - scan.originalPlan.logicalLink.foreach(originalScan.setLogicalLink) - originalScan - case scan: CometIcebergNativeScanExec => - // Iceberg's native scan has the same split between live and frozen filters. - // serializedPartitionData rebuilds originalPlan from runtimeFilters before - // execution, but transition reversion skips that path and executes the restored - // BatchScanExec directly. Carry the executable DPP filters across here as well. - val originalScan = scan.originalPlan.copy(runtimeFilters = scan.runtimeFilters) - scan.originalPlan.logicalLink.foreach(originalScan.setLogicalLink) - originalScan - case _ => cometExec.originalPlan - } - originalWithCurrentExpressions.withNewChildren(cometExec.children) - } else { - logWarning( - "Comet plan and original have different child count for " + - s"${cometExec.getClass.getSimpleName}, using originalPlan as-is.") - cometExec.originalPlan - } + if (isStageBoundary(stripped)) { + throw new CometExec.InvalidSparkFallbackException( + "Cannot revert a stage whose stripped root is a stage boundary") + } + val reverted = transformStageUp(stripped) { case cometExec: CometExec => + cometExec.sparkFallback(cometExec.children) } insertTransitions(reverted) } diff --git a/spark/src/main/scala/org/apache/comet/rules/RewriteJoin.scala b/spark/src/main/scala/org/apache/comet/rules/RewriteJoin.scala index 4ffd34451b4..5f543cc84aa 100644 --- a/spark/src/main/scala/org/apache/comet/rules/RewriteJoin.scala +++ b/spark/src/main/scala/org/apache/comet/rules/RewriteJoin.scala @@ -19,13 +19,16 @@ package org.apache.comet.rules +import org.apache.spark.sql.catalyst.expressions.SortOrder import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide, JoinSelectionHelper} import org.apache.spark.sql.catalyst.plans.{ExistenceJoin, LeftSemi} import org.apache.spark.sql.catalyst.plans.logical.Join import org.apache.spark.sql.execution.{SortExec, SparkPlan} import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} +import org.apache.spark.sql.internal.SQLConf -import org.apache.comet.CometSparkSessionExtensions.withFallbackReason +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.{withFallbackReason, withInfo} /** * Adapted from equivalent rule in Apache Gluten. @@ -64,7 +67,69 @@ object RewriteJoin extends JoinSelectionHelper { case _ => plan } - def rewrite(plan: SparkPlan): SparkPlan = plan match { + /** + * Adds a local sort under any operator whose child no longer has the ordering the operator + * requires, using the same check as EnsureRequirements. EnsureRequirements puts no sort above a + * sort-merge join whose output ordering already satisfies its parent, so when rewrite turns + * that join into a hash join and removes its sorts, a kept parent such as another sort-merge + * join on the same key would read unsorted input. Plans that already satisfy their required + * orderings are returned unchanged. + */ + def restoreRequiredOrdering(plan: SparkPlan): SparkPlan = plan.transformUp { case p => + p.withNewChildren(p.children.zip(p.requiredChildOrdering).map { + case (child, required) if !SortOrder.orderingSatisfies(child.outputOrdering, required) => + SortExec(required, global = false, child = child) + case (child, _) => child + }) + } + + /** + * The largest build side the rewrite converts, or None when the limit is switched off. Unset + * means Spark's own rule for choosing a shuffled hash join by size: the broadcast threshold + * times the initial shuffle partition count. A non-positive threshold only disables broadcast + * joins and says nothing about the hash table an executor can hold, so Spark's default + * threshold stands in for it. + */ + private def maxBuildSize(conf: SQLConf): Option[BigInt] = + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.get(conf) match { + case Some(limit) if limit <= 0 => None + case Some(limit) => Some(BigInt(limit)) + case None => + val threshold = conf.autoBroadcastJoinThreshold + val perPartition = + if (threshold > 0) threshold else SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.defaultValue.get + Some(BigInt(perPartition) * conf.numShufflePartitions) + } + + /** + * Why the build side is too large to rewrite, or None when it fits. Like Spark's + * canBuildLocalHashMapBySize, the size must be strictly under the limit. The hash join cannot + * spill, so a build side without statistics is refused as well. Spark reports an unknown size + * as Long.MaxValue, which fails the comparison on its own. + */ + private def buildSideOverLimit( + smj: SortMergeJoinExec, + buildSide: BuildSide, + conf: SQLConf): Option[String] = + maxBuildSize(conf).flatMap { limit => + smj.logicalLink match { + case Some(join: Join) => + val buildSize = buildSide match { + case BuildLeft => join.left.stats.sizeInBytes + case BuildRight => join.right.stats.sizeInBytes + } + if (buildSize >= limit) { + Some( + s"build side size estimate of $buildSize bytes is not under the limit of " + + s"$limit bytes") + } else { + None + } + case _ => Some("no statistics are available for the build side") + } + } + + def rewrite(plan: SparkPlan, conf: SQLConf): SparkPlan = plan match { case smj: SortMergeJoinExec => getSmjBuildSide(smj) match { case Some(BuildRight) @@ -77,15 +142,23 @@ object RewriteJoin extends JoinSelectionHelper { s"BuildRight with ${smj.joinType} is not supported") plan case Some(buildSide) => - ShuffledHashJoinExec( - smj.leftKeys, - smj.rightKeys, - smj.joinType, - buildSide, - smj.condition, - removeSort(smj.left), - removeSort(smj.right), - smj.isSkewJoin) + buildSideOverLimit(smj, buildSide, conf) match { + case Some(reason) => + // The sort-merge join still runs natively, so this is information for extended + // explain rather than a fallback reason. + withInfo(smj, s"Cannot rewrite SortMergeJoin to HashJoin: $reason") + plan + case None => + ShuffledHashJoinExec( + smj.leftKeys, + smj.rightKeys, + smj.joinType, + buildSide, + smj.condition, + removeSort(smj.left), + removeSort(smj.right), + smj.isSkewJoin) + } case _ => plan } case _ => plan diff --git a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala index 58caedcff14..e07553e2c14 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala @@ -22,7 +22,7 @@ package org.apache.comet.serde import scala.util.control.NonFatal import org.apache.spark.SparkEnv -import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSeq, BindReferences, Expression, Literal, RuntimeReplaceable, ScalaUDF} +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSeq, BindReferences, Expression, ExpressionSet, If, IsNull, KnownNotNull, Literal, Or, RuntimeReplaceable, ScalaUDF} import org.apache.spark.sql.types.BinaryType import org.apache.comet.CometConf @@ -58,6 +58,64 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { override def convert(expr: ScalaUDF, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = emitJvmCodegenDispatch(expr, inputs, binding) + /** + * Emits Spark's null guard around a Scala UDF, together with the UDF, as one dispatcher call. + * Returns `None` when `expr` is not that guard, when an expression in it is disabled, or when + * the dispatcher cannot take it, and records nothing on `expr` in that case: the caller + * converts `expr` as usual, which sends the UDF to the dispatcher on its own, and the UDF + * records its own fallback reason if it cannot go. + * + * Spark's `HandleNullInputsForUDF` rule wraps a UDF with a primitive parameter over a nullable + * input as `if (isnull(a) or isnull(b)) null else f(knownnotnull(a), knownnotnull(b))`, one + * `IsNull` per such argument. Converted natively, the `If` is a `CASE` that splits the batch by + * the predicate and merges the two branches back together, which costs more than the call it + * protects. In the kernel it is a branch per row. + * + * Only that shape matches: the predicate must be the one the rule builds from the UDF's guarded + * arguments, each checked once. The kernel then runs no expression it would not run for the UDF + * alone, apart from the guard's own `if`, `or` and `isnull`, and it runs Spark's code for all + * of them. In a filter or join condition, and in the predicate of a conditional, + * `ReplaceNullWithFalseInPredicate` replaces the null with `false`, so that form matches too. + */ + def emitNullGuardDispatch(expr: If, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = + if (isNullGuard(expr) && isEnabled(expr)) { + tryJvmCodegenDispatch(expr, inputs, binding).toOption + } else { + None + } + + /** + * Whether converting `guard` natively would pass every `spark.comet.expression..enabled` + * check it makes: on the UDF, and on each node of the predicate and the true branch, the + * guarded arguments included. The dispatcher makes none of these checks, so a guard with a + * disabled expression stays native, where the conversion records the fallback reason. The UDF's + * arguments are not checked, because the dispatcher takes them unchecked with or without the + * guard. + */ + private def isEnabled(guard: If): Boolean = { + def disabled(e: Expression): Boolean = + QueryPlanSerde.exprSerdeMap.get(e.getClass).exists { serde => + !CometConf.isExprEnabled( + serde.asInstanceOf[CometExpressionSerde[Expression]].getExprConfigName(e)) + } + !disabled(guard.falseValue) && !guard.predicate.exists(disabled) && + !guard.trueValue.exists(disabled) + } + + private def isNullGuard(expr: If): Boolean = expr match { + case If(predicate, Literal(null, _) | Literal.FalseLiteral, udf: ScalaUDF) => + val guarded = udf.inputPrimitives.zip(udf.children).collect { + case (true, KnownNotNull(arg)) => arg + } + // The optimizer drops a repeated `isnull`, as in `f(a, a)`, so check each argument once. + // `semanticEquals` is false for a nondeterministic argument, so that guard stays native. + ExpressionSet(guarded).toSeq + .map(IsNull(_)) + .reduceLeftOption[Expression](Or) + .exists(_.semanticEquals(predicate)) + case _ => false + } + /** * Bind `expr`, closure-serialize it, and emit a `JvmScalarUdf` proto routed through * [[CometScalaUDFCodegen]] so that native execution evaluates the expression inside the @@ -73,14 +131,28 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { def emitJvmCodegenDispatch( expr: Expression, inputs: Seq[Attribute], - binding: Boolean): Option[Expr] = { + binding: Boolean): Option[Expr] = + tryJvmCodegenDispatch(expr, inputs, binding) match { + case Right(proto) => Some(proto) + case Left(reason) => + withFallbackReason(expr, reason) + None + } + + /** + * [[emitJvmCodegenDispatch]] without the fallback tagging: returns the reason the dispatcher + * cannot take `expr` rather than recording it on `expr`, for a caller that can still convert + * `expr` another way. + */ + private def tryJvmCodegenDispatch( + expr: Expression, + inputs: Seq[Attribute], + binding: Boolean): Either[String, Expr] = { val exprName = CometExplainInfo.exprDisplayName(expr) if (!CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.get()) { - withFallbackReason( - expr, + return Left( s"$exprName: ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=false; expression has " + "no native path so the plan falls back to Spark") - return None } // `RuntimeReplaceable` expressions (e.g. Spark 4's `StructsToJson`) have a `doGenCode` that @@ -97,8 +169,7 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { // under a dispatched `map(...)` would run in it too. Check the whole tree. CometInvokeTargets.declineReason(target) match { case Some(reason) => - withFallbackReason(expr, s"$exprName: codegen dispatch: $reason") - return None + return Left(s"$exprName: codegen dispatch: $reason") case None => } @@ -107,17 +178,15 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { val attrs = target.collect { case a: AttributeReference => a }.distinct val boundExpr = BindReferences.bindReference(target, AttributeSeq(attrs)) - // Gate at plan time. Surface the reason via withFallbackReason rather than crashing Janino - // at execute. + // Gate at plan time. Surface the reason as a fallback rather than crashing Janino at execute. CometBatchKernelCodegen.canHandle(boundExpr) match { case Some(reason) => - withFallbackReason(expr, s"$exprName: $reason") - return None + return Left(s"$exprName: $reason") case None => } // Serialize via Spark's closure serializer: respects the task context classloader (so user - // UDF jars are visible) and matches Spark's wire format. The bytes become arg 0 of the + // UDF jars are visible) and matches Spark's wire format. The bytes become arg 1 of the // JvmScalarUdf proto and self-describe the expression so this works in cluster mode without // executor-side driver registry state. // @@ -139,37 +208,34 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { // unserializable object graph as one of several exception types depending on where in // the graph it trips, and a custom `writeObject` can throw anything. case NonFatal(e) => - withFallbackReason( - expr, + return Left( s"$exprName: codegen dispatch: expression could not be closure-serialized " + s"(${e.getClass.getSimpleName}: ${e.getMessage})") - return None } - val exprArg = exprToProtoInternal(Literal(bytes, BinaryType), inputs, binding).getOrElse { - withFallbackReason( - expr, - s"$exprName: codegen dispatch: could not serialize closure-serialized bound " + - "expression payload") - return None + // Arg 0 is a digest of the bytes. The dispatcher finds each batch's kernel by it, so it reads + // the bytes only to compile on a cache miss instead of copying and hashing them for every + // batch. + val exprArgs = Seq(CometScalaUDFCodegen.digest(bytes), bytes).map { payload => + exprToProtoInternal(Literal(payload, BinaryType), inputs, binding).getOrElse { + return Left( + s"$exprName: codegen dispatch: could not serialize closure-serialized bound " + + "expression payload") + } } val dataArgs = attrs.map { a => exprToProtoInternal(a, inputs, binding).getOrElse { - withFallbackReason(expr, s"$exprName: codegen dispatch: could not serialize data arg $a") - return None + return Left(s"$exprName: codegen dispatch: could not serialize data arg $a") } } val returnTypeProto = serializeDataType(expr.dataType).getOrElse { - withFallbackReason( - expr, - s"$exprName: codegen dispatch: unsupported return type ${expr.dataType}") - return None + return Left(s"$exprName: codegen dispatch: unsupported return type ${expr.dataType}") } val udfBuilder = ExprOuterClass.JvmScalarUdf .newBuilder() .setClassName(classOf[CometScalaUDFCodegen].getName) - .addArgs(exprArg) + exprArgs.foreach(udfBuilder.addArgs) dataArgs.foreach(udfBuilder.addArgs) udfBuilder .setReturnType(returnTypeProto) @@ -194,7 +260,7 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { val _ = withCodegenDispatchExpr(expr, CometExplainInfo.exprDisplayName(node)) case _ => } - Some( + Right( ExprOuterClass.Expr .newBuilder() .setJvmScalarUdf(udfBuilder.build()) diff --git a/spark/src/main/scala/org/apache/comet/serde/CometSortOrder.scala b/spark/src/main/scala/org/apache/comet/serde/CometSortOrder.scala index 1241c0bad38..c70d34dd202 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometSortOrder.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometSortOrder.scala @@ -22,7 +22,6 @@ package org.apache.comet.serde import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, Descending, NullsFirst, NullsLast, SortOrder} import org.apache.spark.sql.types.{ArrayType, DataType, StructType} -import org.apache.comet.CometConf import org.apache.comet.serde.QueryPlanSerde.exprToProtoInternal /** @@ -41,9 +40,8 @@ import org.apache.comet.serde.QueryPlanSerde.exprToProtoInternal * null orders, `ASC NULLS FIRST` and `DESC NULLS LAST`, place it where Spark does. * https://github.com/apache/datafusion-comet/issues/6476 * - * A `RANGE` window frame orders such a null above every other value, so its frames can differ - * from Spark too. Strict mode keeps the fallback it applied to every nested float key before they - * were normalized, for the keys that can hit this. + * A `RANGE` window frame orders such a null above every other value, so `CometWindowExec` + * declines a `RANGE` frame that has to find a row's peers over such a key. * https://github.com/apache/datafusion-comet/issues/6477 */ object CometSortOrder extends CometExpressionSerde[SortOrder] { @@ -62,26 +60,14 @@ object CometSortOrder extends CometExpressionSerde[SortOrder] { "places it by the null order " + "([#6476](https://github.com/apache/datafusion-comet/issues/6476))." - private val nullableNestedFloatingPointSort = - "Sorting on floating-point values nested in an array or struct that can hold a null " + - "element or field" - - override def getIncompatibleReasons(): Seq[String] = Seq( - nestedNullOrderReason, - s"$nullableNestedFloatingPointSort is not 100% compatible with Spark when " + - s"`${CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key}=true`") + override def getIncompatibleReasons(): Seq[String] = Seq(nestedNullOrderReason) override def getSupportLevel(expr: SortOrder): SupportLevel = { - val dataType = expr.child.dataType - if (!canHoldNestedNull(dataType)) { - Compatible() - } else if (expr.nullOrdering != expr.direction.defaultNullOrdering) { + if (canHoldNestedNull(expr.child.dataType) && + expr.nullOrdering != expr.direction.defaultNullOrdering) { Incompatible(Some(nestedNullOrderReason)) } else { - SupportLevel - .strictFloatingPointReason(dataType, nullableNestedFloatingPointSort) - .map(reason => Incompatible(Some(reason))) - .getOrElse(Compatible()) + Compatible() } } diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala b/spark/src/main/scala/org/apache/comet/serde/arrays.scala index b2a75f84782..67ab955296d 100644 --- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala +++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala @@ -115,25 +115,58 @@ object CometArrayAppend extends CometExpressionSerde[ArrayAppend] with ArraysBas object CometArrayContains extends CometExpressionSerde[ArrayContains] + with CometTypeShim with CodegenDispatchFallback { - private val floatingPointReason: String = + override def hasConditionalNativeDefault: Boolean = true + + private val nestedFloatReason: String = "Spark compares array elements with ordering.equiv, so -0.0 matches +0.0 and all NaNs match " + - "each other; Comet's native array_contains compares the raw Arrow values bitwise" - - override def getIncompatibleReasons(): Seq[String] = Seq(floatingPointReason) - - override def getSupportLevel(expr: ArrayContains): SupportLevel = expr.left.dataType match { - // Native array_contains compares floating-point elements bitwise, disagreeing with Spark for - // -0.0/+0.0 and NaN. Report Incompatible (not Unsupported) for float/double element types (at - // any nesting level) so the expression routes through the JVM codegen dispatcher (Spark's own - // doGenCode) and stays native + Spark-exact under the default config, while non-float arrays - // keep the fast native kernel. Under allowIncompatible=true the native kernel is used - // as before. - case ArrayType(elementType, _) - if SupportLevel.containsType(elementType, classOf[FloatType], classOf[DoubleType]) => - Incompatible(Some(floatingPointReason)) - case _ => Compatible() + "each other; Comet's native array_contains applies that only to flat FLOAT and DOUBLE " + + "elements, and nested float elements stay on Spark's comparison" + + private val eagerEvalReason: String = + "array_contains evaluates its value eagerly, while Spark skips it when the array is null" + + private val collationReason: String = + ArrayElementEqualitySupport.collationReason("array_contains") + + override def getIncompatibleReasons(): Seq[String] = + Seq(nestedFloatReason, eagerEvalReason, collationReason) + + private def isFlatFloat(elementType: DataType): Boolean = elementType match { + case FloatType | DoubleType => true + case _ => false + } + + // Whether evaluating `expr` earlier than Spark would is unobservable: a literal or column read + // cannot throw. See CometArrayJoin.orderInsensitive. + private def orderInsensitive(expr: Expression): Boolean = expr match { + case _: Literal | _: Attribute | _: BoundReference => true + case _ => false + } + + override def getSupportLevel(expr: ArrayContains): SupportLevel = { + val elementType = expr.left.dataType.asInstanceOf[ArrayType].elementType + if (hasNonDefaultStringCollation(elementType)) { + // See ArrayElementEqualitySupport: bytewise string comparison ignores the collation. + Incompatible(Some(collationReason)) + } else if (isFlatFloat(elementType)) { + // Spark skips the value when the array is null, so a value that can throw (an ANSI cast, + // say) must not run for a null array. The native kernel evaluates both arguments first. + if (expr.left.nullable && !orderInsensitive(expr.right)) { + Incompatible(Some(eagerEvalReason)) + } else { + Compatible() + } + } else if (SupportLevel.containsType(elementType, classOf[FloatType], classOf[DoubleType])) { + // Float leaves inside arrays or structs stay on the codegen dispatcher: the per-element + // comparator is slower than Spark's generated code there, and it compares any string field + // bytewise, ignoring its collation. + Incompatible(Some(nestedFloatReason)) + } else { + Compatible() + } } override def convert( @@ -142,8 +175,17 @@ object CometArrayContains binding: Boolean): Option[ExprOuterClass.Expr] = { val arrayExprProto = exprToProtoInternal(expr.children.head, inputs, binding) val keyExprProto = exprToProtoInternal(expr.children(1), inputs, binding) - - scalarFunctionExprToProto("array_contains", arrayExprProto, keyExprProto) + // datafusion-spark's array_contains compares floats by their bits. spark_array_contains + // compares them as Spark does, with -0.0 equal to 0.0 and all NaNs equal, at any depth, so + // it also serves nested float elements under allowIncompatible. + val elementType = expr.left.dataType.asInstanceOf[ArrayType].elementType + val function = + if (SupportLevel.containsType(elementType, classOf[FloatType], classOf[DoubleType])) { + "spark_array_contains" + } else { + "array_contains" + } + scalarFunctionExprToProto(function, arrayExprProto, keyExprProto) } } @@ -309,7 +351,19 @@ object CometArrayMin extends CometExpressionSerde[ArrayMin] with CodegenDispatch } } -object CometArraysOverlap extends CometExpressionSerde[ArraysOverlap] { +object CometArraysOverlap + extends CometExpressionSerde[ArraysOverlap] + with CodegenDispatchFallback { + + override def hasConditionalNativeDefault: Boolean = true + + override def getIncompatibleReasons(): Seq[String] = + Seq(ArrayElementEqualitySupport.collationReason("arrays_overlap")) + + // Both inputs share one element type, so checking the left side is enough. + override def getSupportLevel(expr: ArraysOverlap): SupportLevel = + ArrayElementEqualitySupport.getSupportLevel("arrays_overlap", expr.left.dataType) + override def convert( expr: ArraysOverlap, inputs: Seq[Attribute], @@ -538,6 +592,30 @@ object CometSlice extends CometExpressionSerde[Slice] { } } +/** + * Support level for native array kernels that compare elements for equality (membership, overlap, + * dedup). They compare strings by raw bytes, so a non-UTF8_BINARY collation (for example + * UTF8_LCASE, where 'a' equals 'A') gives wrong answers, and these cases report Incompatible. + * Serdes with a boolean result mix in CodegenDispatchFallback, so the JVM codegen dispatcher runs + * Spark's collation-aware comparison. Array-valued serdes fall back instead (see + * ArraySetSupport). + */ +private object ArrayElementEqualitySupport extends CometTypeShim { + def collationReason(name: String): String = + "Spark compares non-UTF8_BINARY collated string elements under their collation, while " + + s"Comet's native $name compares raw bytes" + + def collationSupportLevel(name: String, arrayType: DataType): Option[SupportLevel] = + if (hasNonDefaultStringCollation(arrayType)) { + Some(Incompatible(Some(collationReason(name)))) + } else { + None + } + + def getSupportLevel(name: String, arrayType: DataType): SupportLevel = + collationSupportLevel(name, arrayType).getOrElse(Compatible()) +} + private[comet] object ArraySetSupport { val floatingPointReason: String = "Floating-point elements match Spark's signed-zero and NaN semantics natively only on " + @@ -560,22 +638,33 @@ private[comet] object ArraySetSupport { Compatible() } } + + /** Collated string elements are also Incompatible; see ArrayElementEqualitySupport. */ + def supportLevel(name: String, dataType: DataType): SupportLevel = + ArrayElementEqualitySupport + .collationSupportLevel(name, dataType) + .getOrElse(supportLevel(dataType)) + + def incompatibleReasons(name: String): Seq[String] = + Seq(floatingPointReason, ArrayElementEqualitySupport.collationReason(name)) } // Use projection fallback to avoid codegen dispatch overhead for array-valued results. // The native implementation remains available through opt-in. object CometArrayDistinct extends CometScalarFunction[ArrayDistinct]("array_distinct") { - override def getIncompatibleReasons(): Seq[String] = Seq(ArraySetSupport.floatingPointReason) + override def getIncompatibleReasons(): Seq[String] = + ArraySetSupport.incompatibleReasons("array_distinct") override def getSupportLevel(expr: ArrayDistinct): SupportLevel = - ArraySetSupport.supportLevel(expr.dataType) + ArraySetSupport.supportLevel("array_distinct", expr.dataType) } object CometArrayUnion extends CometExpressionSerde[ArrayUnion] { - override def getIncompatibleReasons(): Seq[String] = Seq(ArraySetSupport.floatingPointReason) + override def getIncompatibleReasons(): Seq[String] = + ArraySetSupport.incompatibleReasons("array_union") override def getSupportLevel(expr: ArrayUnion): SupportLevel = - ArraySetSupport.supportLevel(expr.dataType) + ArraySetSupport.supportLevel("array_union", expr.dataType) override def convert( expr: ArrayUnion, diff --git a/spark/src/main/scala/org/apache/comet/serde/conditional.scala b/spark/src/main/scala/org/apache/comet/serde/conditional.scala index df94be8c93c..6728c9fcf07 100644 --- a/spark/src/main/scala/org/apache/comet/serde/conditional.scala +++ b/spark/src/main/scala/org/apache/comet/serde/conditional.scala @@ -27,6 +27,15 @@ import org.apache.comet.serde.QueryPlanSerde.exprToProtoInternal object CometIf extends CometExpressionSerde[If] { override def convert( + expr: If, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = + // Spark's null guard around a Scala UDF goes to the dispatcher with the UDF when it can. + CometScalaUDF + .emitNullGuardDispatch(expr, inputs, binding) + .orElse(convertNative(expr, inputs, binding)) + + private def convertNative( expr: If, inputs: Seq[Attribute], binding: Boolean): Option[ExprOuterClass.Expr] = { diff --git a/spark/src/main/scala/org/apache/comet/serde/datetime.scala b/spark/src/main/scala/org/apache/comet/serde/datetime.scala index bf27ce2c552..8b3ddab5085 100644 --- a/spark/src/main/scala/org/apache/comet/serde/datetime.scala +++ b/spark/src/main/scala/org/apache/comet/serde/datetime.scala @@ -22,11 +22,11 @@ package org.apache.comet.serde import java.util.Locale import org.apache.spark.sql.catalyst.expressions.{AddMonths, Attribute, Cast, ConvertTimezone, DateAdd, DateDiff, DateFormatClass, DateFromUnixDate, DateSub, DayOfMonth, DayOfWeek, DayOfYear, Days, DivideDTInterval, Expression, FromUTCTimestamp, GetDateField, GetTimestamp, Hour, Hours, LastDay, Literal, MakeDate, MakeDTInterval, MakeInterval, MakeTimestamp, MakeYMInterval, MicrosToTimestamp, MillisToTimestamp, Minute, Month, MonthsBetween, MultiplyDTInterval, MultiplyYMInterval, NextDay, PreciseTimestampConversion, Quarter, Second, SecondsToTimestamp, TimestampAdd, TimestampDiff, ToUnixTimestamp, ToUTCTimestamp, TruncDate, TruncTimestamp, UnixDate, UnixMicros, UnixMillis, UnixSeconds, UnixTimestamp, WeekDay, WeekOfYear, Year} -import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{CalendarIntervalType, DataType, DateType, DoubleType, FloatType, IntegerType, LongType, StringType, TimestampNTZType, TimestampType} import org.apache.spark.unsafe.types.UTF8String import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.isSpark42Plus import org.apache.comet.expressions.{CometCast, CometEvalMode} import org.apache.comet.serde.CometGetDateField.CometGetDateField import org.apache.comet.serde.ExprOuterClass.Expr @@ -69,7 +69,6 @@ trait CometExprGetDateField[T <: GetDateField] { .setChild(e) .setDatatype(serializeDataType(IntegerType).get) .setEvalMode(ExprOuterClass.EvalMode.LEGACY) - .setAllowIncompat(false) .build()) .build() }) @@ -503,7 +502,6 @@ object CometUnixDate extends CometExpressionSerde[UnixDate] { .setChild(child) .setDatatype(serializeDataType(IntegerType).get) .setEvalMode(ExprOuterClass.EvalMode.LEGACY) - .setAllowIncompat(false) .build()) .build() } @@ -644,6 +642,10 @@ object CometTruncTimestamp builder.setChild(childExpr.get) builder.setFormat(formatExpr.get) builder.setTimezone(timeZone.get) + // Spark 4.2.0 and later (SPARK-56663) raise `long overflow` when SECOND or MILLISECOND + // truncation falls below the smallest timestamp; earlier versions wrap the Long + // subtraction. The native side has to match whichever Spark we are running against. + builder.setWrapSecondMillisecondOverflow(!isSpark42Plus) Some( ExprOuterClass.Expr @@ -822,8 +824,14 @@ object CometHours extends CometExpressionSerde[Hours] { * For DateType: dates are internally stored as days since epoch, so this is a simple cast to * integer (same as CometUnixDate). * - * For TimestampType: uses a timezone-aware Cast(Timestamp to Date) followed by Cast(Date to Int). - * The first cast respects the session timezone to correctly determine the date boundary. + * For TimestampType: uses Cast(Timestamp to Date) in UTC followed by Cast(Date to Int). Spark + * cannot evaluate `days` itself, so this counts days in UTC like [[CometHours]] and Iceberg's + * `days` transform, rather than in the session timezone. + * + * The cast is a true floor, as [[CometHours]] is. Iceberg's `DateTimeUtil.microsToDays` differs + * in one case: it puts a pre-1970 timestamp exactly 999999 microseconds past midnight, such as + * `1969-01-01 00:00:00.999999`, in the previous day. Comet's native kernel for Iceberg's own + * `days` function matches Iceberg there, and this cast does not. */ object CometDays extends CometExpressionSerde[Days] { @@ -846,9 +854,8 @@ object CometDays extends CometExpressionSerde[Days] { val dateExprOpt = expr.child.dataType match { case DateType => childExpr case TimestampType => - val timezone = SQLConf.get.sessionLocalTimeZone childExpr.flatMap { child => - CometCast.castToProto(expr, Some(timezone), DateType, child, CometEvalMode.LEGACY) + CometCast.castToProto(expr, Some("UTC"), DateType, child, CometEvalMode.LEGACY) } case _ => None } @@ -863,7 +870,6 @@ object CometDays extends CometExpressionSerde[Days] { .setChild(dateExpr) .setDatatype(serializeDataType(IntegerType).get) .setEvalMode(ExprOuterClass.EvalMode.LEGACY) - .setAllowIncompat(false) .build()) .build() } diff --git a/spark/src/main/scala/org/apache/comet/serde/literals.scala b/spark/src/main/scala/org/apache/comet/serde/literals.scala index 621dd99ebe8..ef6fa1fb466 100644 --- a/spark/src/main/scala/org/apache/comet/serde/literals.scala +++ b/spark/src/main/scala/org/apache/comet/serde/literals.scala @@ -113,7 +113,7 @@ object CometLiteral extends CometExpressionSerde[Literal] with CometTypeShim wit case arr: ArrayType => val listLiteralBuilder: ListLiteral.Builder = - makeListLiteral(value.asInstanceOf[ArrayData].array, arr) + makeListLiteral(value.asInstanceOf[ArrayData].toArray[Any](arr.elementType), arr) exprBuilder.setListVal(listLiteralBuilder.build()) exprBuilder.setDatatype(serializeDataType(dataType).get) case dt => @@ -225,7 +225,7 @@ object CometLiteral extends CometExpressionSerde[Literal] with CometTypeShim wit array.foreach(v => { val casted = v.asInstanceOf[ArrayData] listLiteralBuilder.addListValues(if (casted != null) { - makeListLiteral(casted.array, a) + makeListLiteral(casted.toArray[Any](a.elementType), a) } else ListLiteral.newBuilder()) listLiteralBuilder.addNullMask(casted != null) }) diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala index 6af2cc18ba6..545cbc774ae 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometDataWritingCommand.scala @@ -19,17 +19,21 @@ package org.apache.comet.serde.operator +import java.util.UUID + import scala.jdk.CollectionConverters._ -import org.apache.spark.SparkException -import org.apache.spark.sql.comet.{CometEmptyRelationExec, CometNativeExec, CometNativeWriteExec, CometScanWrapper} -import org.apache.spark.sql.execution.SparkPlan -import org.apache.spark.sql.execution.adaptive.QueryStageExec +import org.apache.hadoop.mapreduce.Job +import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat +import org.apache.spark.internal.io.FileCommitProtocol +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap +import org.apache.spark.sql.comet.{CometNativeExec, CometNativeWriteExec} import org.apache.spark.sql.execution.command.DataWritingCommandExec import org.apache.spark.sql.execution.datasources.{InsertIntoHadoopFsRelationCommand, WriteFilesExec} import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat -import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.util.SerializableConfiguration import org.apache.comet.{CometConf, ConfigEntry} import org.apache.comet.CometSparkSessionExtensions.withFallbackReason @@ -55,25 +59,18 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec case cmd: InsertIntoHadoopFsRelationCommand => cmd.fileFormat match { case _: ParquetFileFormat => - // AQE can replace the write input with a zero-partition empty relation. Keep - // Spark's writer, which creates an empty task to preserve the output file schema. - // The native writer only maps existing partitions; see #5303. This guard is - // conservative: an empty relation below an exchange can have nonzero partitions - // at the write input. Revisit the guard when native empty-file handling is fixed. - if (hasEmptyRelationInput(op.child)) { - return Unsupported(Some( - "Parquet writes with empty-relation inputs require Spark's empty-file handling")) - } - if (!cmd.outputPath.toString.startsWith("file:") && !cmd.outputPath.toString .startsWith("hdfs:")) { return Unsupported(Some("Supported output filesystems: local, HDFS")) } + // Spark 3.x HadoopMapReduceCommitProtocol.getFilename hardcodes "part"; + // mapreduce.output.basename is only honored on Spark 4.0+. Custom committer paths + // are still checked by checkNativeWriteDestination before native execution. NativeWriteUtils - // This writer names its own files `part--.parquet`, so the - // prefix is fixed rather than read from `mapreduce.output.basename`. - .escapedHdfsDestination(cmd.outputPath.toString, "part") + .escapedHdfsDestination( + cmd.outputPath.toString, + NativeWriteUtils.DEFAULT_BASE_OUTPUT_NAME) .foreach(reason => return Unsupported(Some(reason))) if (cmd.bucketSpec.isDefined) { @@ -98,14 +95,6 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec } } - private def hasEmptyRelationInput(plan: SparkPlan): Boolean = plan match { - case _: CometEmptyRelationExec => true - case wrapper: CometScanWrapper => hasEmptyRelationInput(wrapper.originalPlan) - case stage: QueryStageExec => hasEmptyRelationInput(stage.plan) - case reused: ReusedExchangeExec => hasEmptyRelationInput(reused.child) - case _ => plan.children.exists(hasEmptyRelationInput) - } - override def convert( op: DataWritingCommandExec, builder: Operator.Builder, @@ -123,6 +112,8 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec val outputPath = cmd.outputPath.toString + // Planning-time value only. CometNativeWriteExec replaces it per task with the codec that + // names the file. val plannedCodec = NativeWriteUtils.parseCompressionCodec(cmd.options) val codec = NativeWriteUtils.protoCompressionCodec(plannedCodec) match { case Some(codec) => codec @@ -140,8 +131,8 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec cmd.query.schema.fields.toIndexedSeq, Some( op.session.sessionState.conf.getConf(SQLConf.PARQUET_FIELD_ID_WRITE_ENABLED))).asJava) - // Note: work_dir, job_id, and task_attempt_id will be set at execution time - // in CometNativeWriteExec, as they depend on the Spark task context + // CometNativeWriteExec replaces output_path with the committer's exact task filename + // at execution time. // Collect S3/cloud storage configurations val session = op.session @@ -189,29 +180,30 @@ object CometDataWritingCommand extends CometOperatorSerde[DataWritingCommandExec other } - // Create FileCommitProtocol for atomic writes - val jobId = java.util.UUID.randomUUID().toString - val committer = - try { - // Use Spark's SQLHadoopMapReduceCommitProtocol - val committerClass = - classOf[org.apache.spark.sql.execution.datasources.SQLHadoopMapReduceCommitProtocol] - val constructor = - committerClass.getConstructor(classOf[String], classOf[String], classOf[Boolean]) - Some( - constructor - .newInstance( - jobId, - outputPath, - java.lang.Boolean.FALSE // dynamicPartitionOverwrite = false for now - ) - .asInstanceOf[org.apache.spark.internal.io.FileCommitProtocol]) - } catch { - case e: Exception => - throw new SparkException(s"Could not instantiate FileCommitProtocol: ${e.getMessage}") - } - - CometNativeWriteExec(nativeOp, childPlan, outputPath, cmd.mode, committer, jobId) + val session = op.session + val job = Job.getInstance(session.sessionState.newHadoopConfWithOptions(cmd.options)) + job.setOutputKeyClass(classOf[Void]) + job.setOutputValueClass(classOf[InternalRow]) + FileOutputFormat.setOutputPath(job, cmd.outputPath) + val outputWriterFactory = + cmd.fileFormat.prepareWrite(session, job, CaseInsensitiveMap(cmd.options), cmd.query.schema) + + val committer = FileCommitProtocol.instantiate( + session.sessionState.conf.fileCommitProtocolClass, + UUID.randomUUID().toString, + outputPath, + dynamicPartitionOverwrite = false) + job.getConfiguration.set("spark.sql.sources.writeJobUUID", UUID.randomUUID().toString) + + CometNativeWriteExec( + nativeOp, + op, + childPlan, + outputPath, + cmd.mode, + committer, + new SerializableConfiguration(job.getConfiguration), + outputWriterFactory) } } diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala index d7865833776..335f02ff034 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeScan.scala @@ -282,6 +282,19 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit deleteFile: Any): Option[AnyRef] = IcebergReflection.findMethod(clazz, methodName).flatMap(m => Option(m.invoke(deleteFile))) + /** + * Reads equality-delete field IDs for the native serde path. The accessor is required on every + * supported Iceberg DeleteFile API, so a missing method or invocation failure must propagate. A + * null return is distinct: Iceberg uses it for files without equality keys. + */ + private def requiredEqualityFieldIds( + deleteFileClass: Class[_], + deleteFile: Any): java.util.List[Integer] = { + val method = IcebergReflection.getMethod(deleteFileClass, "equalityFieldIds") + val ids = method.invoke(deleteFile).asInstanceOf[java.util.List[Integer]] + if (ids == null) new java.util.ArrayList[Integer]() else ids + } + /** * Extracts delete files from an Iceberg FileScanTask as a list (for deduplication). * @@ -301,97 +314,15 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit val deletes = IcebergReflection.getDeleteFilesFromTask(task, fileScanTaskClass) - deletes.asScala.map { deleteFile => - // The path is the one essential field. A delete file we cannot locate cannot be applied, - // and silently skipping it would leak deleted rows, so treat a missing path as fatal. - val deletePath = IcebergReflection - .extractFileLocation(contentFileClass, deleteFile) - .getOrElse( - throw new RuntimeException( - "Neither location() nor path() is declared on this Iceberg version's " + - "ContentFile -- cannot extract delete file path from FileScanTask")) - - val deleteBuilder = OperatorOuterClass.IcebergDeleteFile.newBuilder() - deleteBuilder.setFilePathIdx(internPath(deletePath)) - - val contentType = - try { - val contentMethod = IcebergReflection.getMethod(deleteFileClass, "content") - val content = contentMethod.invoke(deleteFile) - content.toString match { - case IcebergReflection.ContentTypes.POSITION_DELETES => - IcebergReflection.ContentTypes.POSITION_DELETES - case IcebergReflection.ContentTypes.EQUALITY_DELETES => - IcebergReflection.ContentTypes.EQUALITY_DELETES - case other => other - } - } catch { - case _: Exception => - IcebergReflection.ContentTypes.POSITION_DELETES - } - deleteBuilder.setContentType(contentType) - - // "PARQUET", or "PUFFIN" for a V3 deletion vector. iceberg-rust selects its - // deletion-vector reader on this, and a wrong value reads a Puffin blob as Parquet, so an - // undeterminable format is fatal: by serde time there is no fallback left. - val fileFormat = IcebergReflection - .getFileFormat(contentFileClass, deleteFile) - .getOrElse( - throw new RuntimeException( - "ContentFile.format() is not declared on this Iceberg version -- cannot tell a " + - "deletion vector from a Parquet delete file")) - deleteBuilder.setFileFormat(fileFormat) - - val specId = - try { - val specIdMethod = IcebergReflection.getMethod(deleteFileClass, "specId") - specIdMethod.invoke(deleteFile).asInstanceOf[Int] - } catch { - case _: Exception => 0 - } - deleteBuilder.setPartitionSpecId(specId) - - try { - val equalityIdsMethod = - IcebergReflection.getMethod(deleteFileClass, "equalityFieldIds") - val equalityIds = equalityIdsMethod - .invoke(deleteFile) - .asInstanceOf[java.util.List[Integer]] - deleteBuilder.addAllEqualityIds(equalityIds) - } catch { - case _: Exception => - } - - // Gated on the format, not on the accessors returning a value: Iceberg also sets - // referencedDataFile on file-scoped Parquet position deletes, where iceberg-rust ignores - // it. Forwarding it there would serialize a data-file path per delete file that nothing - // reads, and would invite keying deletion-vector detection on the field instead of on - // fileFormat, which is the only discriminator. - if (fileFormat.equalsIgnoreCase(IcebergReflection.FileFormats.PUFFIN)) { - deletionVectorField(deleteFileClass, "referencedDataFile", deleteFile) - .foreach(p => deleteBuilder.setReferencedDataFile(p.asInstanceOf[String])) - deletionVectorField(deleteFileClass, "contentOffset", deleteFile) - .foreach(o => deleteBuilder.setContentOffset(o.asInstanceOf[java.lang.Long])) - deletionVectorField(deleteFileClass, "contentSizeInBytes", deleteFile) - .foreach(s => deleteBuilder.setContentSizeInBytes(s.asInstanceOf[java.lang.Long])) - } - - // recordCount is declared on ContentFile, so it is present on every supported Iceberg - // version and a lookup failure is a real defect rather than an old-version absence. - // iceberg-rust rejects a deletion vector without it, since it checks the count against - // the cardinality it decodes from the blob. - deleteBuilder.setRecordCount( - IcebergReflection - .getMethod(contentFileClass, "recordCount") - .invoke(deleteFile) - .asInstanceOf[java.lang.Long]) - - // Encrypted delete files carry a plaintext StandardKeyMetadata blob; forward it verbatim. - // Unencrypted delete files leave the field unset. - keyMetadataBytes(keyMetadataMethod, deleteFile).foreach(deleteBuilder.setKeyMetadata) - - deleteBuilder.build() - }.toSeq + deletes.asScala + .map( + serializeDeleteFile( + _, + contentFileClass, + deleteFileClass, + keyMetadataMethod, + internPath)) + .toSeq } catch { case e: Exception => val msg = @@ -402,6 +333,85 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit } } + /** + * Serializes a single Iceberg DeleteFile to protobuf. + * + * Required delete metadata must never fall back to guessed values after native execution has + * been selected. Missing accessors and invocation failures therefore propagate. + */ + private[operator] def serializeDeleteFile( + deleteFile: Any, + contentFileClass: Class[_], + deleteFileClass: Class[_], + keyMetadataMethod: java.lang.reflect.Method, + internPath: String => Int): OperatorOuterClass.IcebergDeleteFile = { + val deletePath = IcebergReflection + .extractFileLocation(contentFileClass, deleteFile) + .getOrElse( + throw new RuntimeException( + "Neither location() nor path() is declared on this Iceberg version's " + + "ContentFile -- cannot extract delete file path from FileScanTask")) + + val deleteBuilder = OperatorOuterClass.IcebergDeleteFile.newBuilder() + deleteBuilder.setFilePathIdx(internPath(deletePath)) + + val contentMethod = IcebergReflection.getMethod(deleteFileClass, "content") + val contentType = contentMethod.invoke(deleteFile).toString + deleteBuilder.setContentType(contentType) + + // "PARQUET", or "PUFFIN" for a V3 deletion vector. iceberg-rust selects its + // deletion-vector reader on this, and a wrong value reads a Puffin blob as Parquet, so an + // undeterminable format is fatal: by serde time there is no fallback left. + val fileFormat = IcebergReflection + .getFileFormat(contentFileClass, deleteFile) + .getOrElse( + throw new RuntimeException( + "ContentFile.format() is not declared on this Iceberg version -- cannot tell a " + + "deletion vector from a Parquet delete file")) + deleteBuilder.setFileFormat(fileFormat) + + val specIdMethod = IcebergReflection.getMethod(deleteFileClass, "specId") + deleteBuilder.setPartitionSpecId(specIdMethod.invoke(deleteFile).asInstanceOf[Int]) + + val equalityFieldIds = requiredEqualityFieldIds(deleteFileClass, deleteFile) + val isEqualityDelete = contentType == IcebergReflection.ContentTypes.EQUALITY_DELETES + if (isEqualityDelete && equalityFieldIds.isEmpty) { + throw new IllegalStateException( + s"Iceberg equality delete file '$deletePath' has no equality field IDs") + } + deleteBuilder.addAllEqualityIds(equalityFieldIds) + + // Gated on the format, not on the accessors returning a value: Iceberg also sets + // referencedDataFile on file-scoped Parquet position deletes, where iceberg-rust ignores + // it. Forwarding it there would serialize a data-file path per delete file that nothing + // reads, and would invite keying deletion-vector detection on the field instead of on + // fileFormat, which is the only discriminator. + if (fileFormat.equalsIgnoreCase(IcebergReflection.FileFormats.PUFFIN)) { + deletionVectorField(deleteFileClass, "referencedDataFile", deleteFile) + .foreach(p => deleteBuilder.setReferencedDataFile(p.asInstanceOf[String])) + deletionVectorField(deleteFileClass, "contentOffset", deleteFile) + .foreach(o => deleteBuilder.setContentOffset(o.asInstanceOf[java.lang.Long])) + deletionVectorField(deleteFileClass, "contentSizeInBytes", deleteFile) + .foreach(s => deleteBuilder.setContentSizeInBytes(s.asInstanceOf[java.lang.Long])) + } + + // recordCount is declared on ContentFile, so it is present on every supported Iceberg + // version and a lookup failure is a real defect rather than an old-version absence. + // iceberg-rust rejects a deletion vector without it, since it checks the count against + // the cardinality it decodes from the blob. + deleteBuilder.setRecordCount( + IcebergReflection + .getMethod(contentFileClass, "recordCount") + .invoke(deleteFile) + .asInstanceOf[java.lang.Long]) + + // Encrypted delete files carry a plaintext StandardKeyMetadata blob; forward it verbatim. + // Unencrypted delete files leave the field unset. + keyMetadataBytes(keyMetadataMethod, deleteFile).foreach(deleteBuilder.setKeyMetadata) + + deleteBuilder.build() + } + /** * Serializes partition spec and data from an Iceberg FileScanTask. * @@ -1134,10 +1144,7 @@ object CometIcebergNativeScan extends CometOperatorSerde[CometBatchScanExec] wit // dropped ones from the table's schema history (mirrors Iceberg-Java's // DeleteFilter.fileProjection). val equalityFieldIds = deletes.asScala.flatMap { df => - IcebergReflection - .getEqualityFieldIds(deleteFileClass, df) - .asScala - .map(_.asInstanceOf[java.lang.Integer].intValue()) + requiredEqualityFieldIds(deleteFileClass, df).asScala.map(_.intValue()) }.toSeq if (equalityFieldIds.nonEmpty) { IcebergReflection diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala index c6741f3e40c..9c708e2d0d1 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala @@ -30,7 +30,7 @@ import org.apache.spark.sql.comet.{CometIcebergWriteExec, CometNativeExec, Icebe import org.apache.comet.{CometConf, ConfigEntry} import org.apache.comet.CometSparkSessionExtensions.withFallbackReason -import org.apache.comet.iceberg.IcebergReflection +import org.apache.comet.iceberg.{IcebergReflection, IcebergStorageSchemes, PositionDeltaWrite, ReplaceDataWrite} import org.apache.comet.objectstore.NativeConfig import org.apache.comet.serde.{CometOperatorSerde, Compatible, OperatorOuterClass, SupportLevel, Unsupported} import org.apache.comet.serde.OperatorOuterClass.Operator @@ -86,12 +86,11 @@ object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] { // `timestamp_ns`, `geometry` or `geography` today, so those are declined in case it learns to. private val UnsupportedWriteTypeIds: Set[String] = Set("UUID", "VARIANT", "UNKNOWN", "TIMESTAMP_NANO", "GEOMETRY", "GEOGRAPHY") - // `oss` is deliberately absent: iceberg-rust has an OSS backend, but Comet does not forward - // `oss.*` catalog properties to it and no functional test covers the path, so an OSS write - // could silently drop endpoint/credential configuration. Fail closed until it is covered. - // `gs` is additionally gated on the resolved FileIO (`requireGcsFileIOForGcsDataLocation`). - private val SupportedStorageSchemes: Set[String] = - Set("file", "memory", "s3", "s3a", "gs") + // Loaded from the native storage factory: `builtin_storage_schemes` in + // `native/core/src/execution/operators/iceberg_common.rs` is the single point of change and + // explains why `oss` is read-only and `memory` write-only. Lazy so constructing the serde does + // not touch the native library. `gs` is additionally gated on the resolved FileIO below. + private lazy val SupportedStorageSchemes: Set[String] = IcebergStorageSchemes.write // Supported schemes whose native backend is local and needs no host. Every other supported // scheme reads its bucket from the URL host (`requireSupportedStorageScheme`). private val LocalStorageSchemes: Set[String] = Set("file", "memory") @@ -159,10 +158,19 @@ object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] { // from spark.hadoop.*. They do not alter an Iceberg data-file write request, so they must not // make every otherwise-clean S3 write ineligible. This also permits explicit overrides, which // are harmless on the write path for the same reason. + // + // Some Spark distributions also set S3A delete, read and committer defaults cluster-wide. The + // native writer does not go through S3A, `fadvise` is a read hint, and Iceberg commits through + // table metadata rather than a Hadoop output committer, so these are harmless here as well. private val IgnoredHadoopS3Keys: Set[String] = Set( "fs.s3a.downgrade.syncable.exceptions", "fs.s3a.vectored.read.max.merged.size", - "fs.s3a.vectored.read.min.seek.size") + "fs.s3a.vectored.read.min.seek.size", + "fs.s3a.bulk.delete.page.size", + "fs.s3a.experimental.input.fadvise", + "fs.s3a.committer.magic.enabled", + "fs.s3a.committer.name", + "fs.s3a.committer.threads") // Audited against the pinned iceberg-rust S3 parser // (`iceberg/src/io/storage/config/s3.rs` and `storage/opendal/src/s3.rs`). Do not broaden this @@ -292,6 +300,12 @@ object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] { } private def checkTriggers(op: IcebergWriteExec): Option[String] = { + op.dispatch match { + case PositionDeltaWrite(_) => + return Some("Iceberg WriteDelta executes through the JVM DeltaWriter") + case _ => + } + val batchWrite = op.batchWrite if (!IcebergReflection.isIcebergBatchWrite(batchWrite)) { return Some(s"not an Iceberg SparkWrite: ${batchWrite.getClass.getName}") @@ -935,13 +949,12 @@ object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] { * giving it the wider row (e.g. 6 columns when the schema has 3) and * `decorate_batch_with_field_ids` rejects the batch. * - * For Spark 3.4 / 3.5 the strategy shim returns `None` for `replaceDataDispatch` and the - * upstream plan already projects to the data columns -- no extra projection needed. For 4.x we - * splice a `Projection` proto between our `IcebergWrite` op and the FFI `Scan`, selecting the - * upstream attributes whose names match the Iceberg schema's columns. The JVM-side child stays - * at the original wide output, so its `executeColumnar()` still emits the wide batches the FFI - * scan declares; the projection then strips them inside the native runtime before the writer - * sees the data. + * Plain writes already present only data columns, so no extra projection is needed. For 4.x + * ReplaceData we splice a `Projection` proto between our `IcebergWrite` op and the FFI `Scan`, + * selecting the upstream attributes whose names match the Iceberg schema's columns. The + * JVM-side child stays at the original wide output, so its `executeColumnar()` still emits the + * wide batches the FFI scan declares; the projection then strips them inside the native runtime + * before the writer sees the data. */ private def dropNonDataColumns( op: IcebergWriteExec, @@ -950,7 +963,10 @@ object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] { // write schema has no row lineage columns: when it has, Iceberg's writer reads their values // from the metadata columns (`ExtractRowLineage`), which this projection discards. // `requireNoMetadataColumns` declines those writes. - if (op.replaceDataDispatch.isEmpty) return Some(scan) + op.dispatch match { + case ReplaceDataWrite(_) => + case _ => return Some(scan) + } val sparkWrite = IcebergReflection.getOuterSparkWrite(op.batchWrite).getOrElse { withFallbackReason(op, "Could not unwrap outer SparkWrite for ReplaceData projection") @@ -1024,7 +1040,9 @@ object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] { "Native Iceberg write conversion: SparkWrite.outputSpecId reflection failed")) CometIcebergWriteExec( nativeOp, + op, op.child, + op.output, op.batchWrite, table.asInstanceOf[AnyRef], outputSpecId) diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala b/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala index 5246099fdcb..a0ebdea520f 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/NativeWriteUtils.scala @@ -141,12 +141,13 @@ object NativeWriteUtils { * byte for byte (see [[hdfsPathDivergence]] for why it may not): * * - the destination directory, and - * - `fileNamePrefix`, the basename every file name is built from. On Spark 4.0+ that is - * `mapreduce.output.basename`, which `HadoopMapReduceCommitProtocol.getFilename` - * interpolates into `--`; on 3.x Comet names the files itself and - * the basename is always the literal `part`. A basename holding `?` or `#` is the dangerous - * one: the native URL parser truncates there, so *every* task writes a file with the same - * truncated name and they overwrite each other during commit. + * - `fileNamePrefix`, the basename every file name is built from. Spark 4.0+ + * `HadoopMapReduceCommitProtocol.getFilename` reads `mapreduce.output.basename` and + * interpolates it into `--`; Spark 3.x hardcodes `part` instead. + * Both native writers take the exact file name from the commit protocol. A basename holding + * `?` or `#` is dangerous on 4.0+: the native URL parser truncates there, so every task + * writes a file with the same truncated name and they overwrite each other during commit. + * Custom committer paths are checked by [[checkNativeWriteDestination]]. * * The basename is checked by running [[hdfsPathDivergence]] over the path it produces rather * than over the name alone. Everything else `getFilename` interpolates -- the split number, the diff --git a/spark/src/main/scala/org/apache/comet/udf/CometNativeUDF.scala b/spark/src/main/scala/org/apache/comet/udf/CometNativeUDF.scala index 5c2a9608307..0feb9605922 100644 --- a/spark/src/main/scala/org/apache/comet/udf/CometNativeUDF.scala +++ b/spark/src/main/scala/org/apache/comet/udf/CometNativeUDF.scala @@ -23,6 +23,7 @@ import org.apache.spark.sql.SparkSession import org.apache.spark.sql.comet.CometUdfErrors import org.apache.spark.sql.types.DataType +import org.apache.comet.serde.QueryPlanSerde import org.apache.comet.shims.ShimSessionFunctionRegistry /** @@ -46,9 +47,10 @@ object CometNativeUDF { /** * Register a single native UDF with an explicit signature. * - * Validates the library on the driver (loads it, confirms a UDF named `name` exists), then - * installs `name` as a temporary function of the session, as `spark.udf.register` would: other - * sessions do not see it, and registering another function under the same name replaces it. + * Checks that Comet can carry the argument and return types natively and validates the library + * on the driver (loads it, confirms a UDF named `name` exists), then installs `name` as a + * temporary function of the session, as `spark.udf.register` would: other sessions do not see + * it, and registering another function under the same name replaces it. * * Executors need no registration: the library path travels with the plan in the * `NativeScalarUdf` proto, and each executor loads the library itself on first use. The path @@ -58,8 +60,9 @@ object CometNativeUDF { * argument types differ, other than in nullability, and inserts no casts, so cast the arguments * in the query instead. * - * Spark cannot evaluate the call itself: if Comet does not take the operator holding it, the - * query fails. + * Spark cannot evaluate the call itself, so a query fails where Spark would have to: in an + * operator Comet does not take, and while planning over local data, in a filter on partition + * columns, or to sample the keys of a global sort. * * `deterministic` must be `true`. Comet plans every imported kernel as immutable, so a * nondeterministic UDF cannot yet be expressed; passing `false` fails here rather than silently @@ -84,6 +87,13 @@ object CometNativeUDF { "eliminated as a common subexpression. " + "See https://github.com/apache/datafusion-comet/issues/5249") } + // With a type Comet cannot represent natively, every call would fall back to Spark, which + // cannot evaluate it. + (inputTypes :+ returnType).find(QueryPlanSerde.serializeDataType(_).isEmpty).foreach { t => + throw new IllegalArgumentException( + s"native UDF '$name': Comet has no native representation for type " + + s"${t.catalogString}, so no call to it could run.") + } validateLibrary(libraryPath, name) ShimSessionFunctionRegistry .functionRegistry(spark) diff --git a/spark/src/main/scala/org/apache/comet/udf/CometUDF.scala b/spark/src/main/scala/org/apache/comet/udf/CometUDF.scala index 6b435c40643..a1d72939914 100644 --- a/spark/src/main/scala/org/apache/comet/udf/CometUDF.scala +++ b/spark/src/main/scala/org/apache/comet/udf/CometUDF.scala @@ -28,6 +28,9 @@ import org.apache.arrow.vector.ValueVector * - Vector arguments arrive at the row count of the current batch. * - Scalar (literal-folded) arguments arrive as length-1 vectors and must be read at index 0. * - The returned vector's length must match `numRows`. + * - Inputs belong to native execution and must not be modified. Each call receives new vectors, + * but a scalar argument's buffers are the same on every call for the life of the expression, + * and a vector argument can share its buffers with the rest of the plan. * * `numRows` mirrors DataFusion's `ScalarFunctionArgs.number_rows` and is the batch row count. * UDFs that always have at least one batch-length input can read length from it and ignore @@ -45,4 +48,23 @@ import org.apache.arrow.vector.ValueVector */ trait CometUDF { def evaluate(inputs: Array[ValueVector], numRows: Int): ValueVector + + /** + * The overload the bridge calls. `partitionIndex` is the index of the partition the calling + * native plan computes, which differs from `TaskContext.partitionId()` under a union, a + * coalesce or a cartesian product. `planId` tells apart the plans one task runs, such as the + * parent partitions of a coalesce. The default ignores both. + */ + def evaluate( + inputs: Array[ValueVector], + numRows: Int, + partitionIndex: Int, + planId: Long): ValueVector = + evaluate(inputs, numRows) + + /** + * Called once native plan `planId` has closed, so an implementation that keeps state per plan + * can drop it before the task ends. The default does nothing. + */ + def releasePlan(planId: Long): Unit = () } diff --git a/spark/src/main/scala/org/apache/comet/udf/CometUdfExceptions.scala b/spark/src/main/scala/org/apache/comet/udf/CometUdfExceptions.scala index eeda2447b1f..44e0554653d 100644 --- a/spark/src/main/scala/org/apache/comet/udf/CometUdfExceptions.scala +++ b/spark/src/main/scala/org/apache/comet/udf/CometUdfExceptions.scala @@ -21,12 +21,11 @@ package org.apache.comet.udf import org.apache.comet.CometRuntimeException -/** - * Thrown when Spark itself evaluates a call to a UDF registered with Comet, which means Comet did - * not take the operator holding it. - */ +/** Thrown when Spark itself evaluates a call to a native UDF, which only Comet can evaluate. */ class CometUdfNotEvaluatedException(name: String) extends CometRuntimeException( s"UDF '$name' is registered with Comet and runs only inside Comet's native execution, but " + - "Spark evaluated it, which means Comet did not take the operator holding the call. The " + - "query's extended explain output gives the reason that operator fell back to Spark.") + "Spark evaluated it. Either Comet did not take the operator holding the call, and the " + + "query's extended explain output gives the reason, or Spark evaluated the call itself " + + "while planning: over local data, in a filter on partition columns, or to sample the " + + "keys of a global sort.") diff --git a/spark/src/main/scala/org/apache/comet/udf/NativeUdfCall.scala b/spark/src/main/scala/org/apache/comet/udf/NativeUdfCall.scala index 22c71039a2d..d1f444cc4ba 100644 --- a/spark/src/main/scala/org/apache/comet/udf/NativeUdfCall.scala +++ b/spark/src/main/scala/org/apache/comet/udf/NativeUdfCall.scala @@ -29,9 +29,10 @@ import org.apache.spark.sql.types.DataType * in the session builds one for each call it resolves. * * Comet replaces it with a `NativeScalarUdf` that runs the function named `name` in the library - * at `libraryPath`. Spark has no way to evaluate it, so `eval` throws. That happens only when - * Comet does not take the operator holding the call, and failing there makes the fallback visible - * rather than silent. + * at `libraryPath`. Spark has no way to evaluate it, so `eval` throws, which makes a fallback + * visible rather than silent. Spark evaluates a call where Comet does not take the operator + * holding it, and while planning in a few places: over local data, in a filter on partition + * columns, and to sample the keys of a global sort. * * `argumentTypes` is the registered signature, one type per child. Spark's analyzer checks each * call against it, disregarding nullability, and inserts no casts. diff --git a/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala b/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala index 1cab882df93..11b2594fdde 100644 --- a/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala +++ b/spark/src/main/scala/org/apache/comet/udf/codegen/CometScalaUDFCodegen.scala @@ -20,6 +20,7 @@ package org.apache.comet.udf.codegen import java.nio.ByteBuffer +import java.security.MessageDigest import java.util.Collections import java.util.concurrent.atomic.AtomicLong @@ -29,7 +30,7 @@ import scala.util.control.NonFatal import org.apache.arrow.vector._ import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector} import org.apache.arrow.vector.types.pojo.Field -import org.apache.spark.{SparkEnv, TaskContext} +import org.apache.spark.SparkEnv import org.apache.spark.internal.Logging import org.apache.spark.sql.catalyst.expressions.Expression import org.apache.spark.sql.comet.util.Utils @@ -41,11 +42,14 @@ import org.apache.comet.udf.CometUDF /** * Arrow-direct codegen dispatcher. For each `(bound expression, input Arrow schema)` pair, - * compiles a specialized [[CometBatchKernel]] on first encounter, initializes it with the task's - * partition index, and caches the live instance. + * compiles a specialized [[CometBatchKernel]] on first encounter, initializes it with the index + * of the partition the calling native plan computes, and caches the live instance. * - * Arg 0 is a `VarBinaryVector` scalar carrying the closure-serialized bound `Expression` bytes; - * args 1..N are the data columns the `BoundReference`s read in ordinal order. + * Arg 0 is a `VarBinaryVector` scalar carrying the [[CometScalaUDFCodegen.digest]] of the + * closure-serialized bound `Expression`; arg 1 is a `VarBinaryVector` scalar carrying the + * serialized bytes themselves; args 2..N are the data columns the `BoundReference`s read in + * ordinal order. A batch finds its kernel by the digest, so the serialized expression, several KB + * for a Scala UDF's closure, is read only to compile a kernel on a cache miss. * * Caching hierarchy, broadest scope on the left: * {{{ @@ -56,16 +60,20 @@ import org.apache.comet.udf.CometUDF * +----------------------------+ +----------------------------+ +----------------------------+ * | Key: generated Java | | Key: task + UDF class | | Key: bound expression + | * | source | | | | input column shapes | - * | Value: compiled Java class | | Value: dispatcher object | | Value: ready-to-run kernel | - * | Scope: JVM, all queries | | Scope: one Spark task | | with state primed | - * | share it | | | | Scope: one Spark task | - * | Owner: Spark | | Owner: Comet | | (lives inside 2) | + * | Value: compiled Java class | | Value: dispatcher object | | (+ native plan if | + * | Scope: JVM, all queries | | Scope: one Spark task | | nondeterministic) | + * | share it | | | | Value: ready-to-run kernel | + * | Owner: Spark | | Owner: Comet | | with state primed | + * | | | | | Scope: one Spark task | + * | | | | | (lives inside 2), | + * | | | | | or one native plan | + * | | | | | if nondeterministic | * | | | | | Owner: Comet | * +----------------------------+ +----------------------------+ +----------------------------+ * }}} * - * Stateful expressions (`Rand`, `MonotonicallyIncreasingID`) advance inside the per-task kernel - * across batches. + * Stateful expressions (`Rand`, `MonotonicallyIncreasingID`) advance inside the per-plan kernel + * across batches. `CometExecIterator.close` drops a plan's kernels through `releasePlan`. * * `evaluate` runs under `this.synchronized` because DataFusion operators like `HashJoinExec` * pipeline build/probe via `OnceAsync` (`tokio::spawn`), so multiple Tokio worker threads can @@ -78,36 +86,57 @@ import org.apache.comet.udf.CometUDF class CometScalaUDFCodegen extends CometUDF with Logging { /** - * Per-task cache keyed on serialized expression bytes plus per-column specs. The deserialized - * `boundExpr` carries mutable state (`NamedLambdaVariable.value` for HOFs, `Rand`'s - * `XORShiftRandom`) that must not be shared across concurrent tasks running the same query; - * keeping the cache per-task gives each task its own copy. Guarded by `this.synchronized`. + * Per-task cache keyed on the serialized expression's digest plus per-column specs. The + * deserialized `boundExpr` carries mutable state (`NamedLambdaVariable.value` for HOFs, + * `Rand`'s `XORShiftRandom`) that must not be shared across concurrent tasks running the same + * query; keeping the cache per-task gives each task its own copy. A nondeterministic kernel is + * seeded from the partition its plan computes, so its key also holds the plan: each plan a task + * runs, such as each parent partition of a coalesce, gets its own, dropped by `releasePlan` + * when the plan closes. Guarded by `this.synchronized`. */ private val kernelCache : mutable.Map[CometScalaUDFCodegen.CacheKey, CometScalaUDFCodegen.CacheEntry] = mutable.HashMap.empty - override def evaluate(inputs: Array[ValueVector], numRows: Int): ValueVector = { + // Kernels shared by every plan stay until the task ends. + override def releasePlan(planId: Long): Unit = this.synchronized { + kernelCache.keys.filter(_.planId == planId).toList.foreach(kernelCache.remove) + } + + /** Plan ids of the cached kernels, `NoPlan` for a shared one. */ + private[comet] def cachedPlanIds: List[Long] = this.synchronized { + kernelCache.keysIterator.map(_.planId).toList.sorted + } + + // Callers that bypass the bridge (unit tests, benchmarks) have no native plan. + override def evaluate(inputs: Array[ValueVector], numRows: Int): ValueVector = + evaluate(inputs, numRows, partitionIndex = 0, planId = CometScalaUDFCodegen.NoPlan) + + override def evaluate( + inputs: Array[ValueVector], + numRows: Int, + partitionIndex: Int, + planId: Long): ValueVector = { require( - inputs.length >= 1, - "CometScalaUDFCodegen requires at least 1 input (serialized expression), " + - s"got ${inputs.length}") - val exprVec = inputs(0).asInstanceOf[VarBinaryVector] + inputs.length >= 2, + "CometScalaUDFCodegen requires at least 2 inputs (expression digest and serialized " + + s"expression), got ${inputs.length}") + val digest = binaryScalar(inputs(0), "expression digest at arg 0") require( - exprVec.getValueCount >= 1 && !exprVec.isNull(0), - "CometScalaUDFCodegen requires non-null serialized expression bytes at arg 0") - val bytes = exprVec.get(0) + digest.length == CometScalaUDFCodegen.DigestLength, + s"CometScalaUDFCodegen requires a ${CometScalaUDFCodegen.DigestLength}-byte expression " + + s"digest at arg 0, got ${digest.length} bytes") // TODO(dict-encoded): kernels assume materialized inputs. Dict-encoded vectors would fail the // cast in `specFor` below. Fix is to materialize at the dispatcher (via // `CDataDictionaryProvider`) or widen `emitTypedGetters` with a dict-index + lookup path. - val numDataCols = inputs.length - 1 + val numDataCols = inputs.length - 2 val dataCols = new Array[ValueVector](numDataCols) val specs = new Array[ArrowColumnSpec](numDataCols) var di = 0 while (di < numDataCols) { - val v = inputs(di + 1) + val v = inputs(di + 2) dataCols(di) = v specs(di) = specFor(v) di += 1 @@ -115,12 +144,12 @@ class CometScalaUDFCodegen extends CometUDF with Logging { val n = numRows val specsSeq = specs.toIndexedSeq - val key = CometScalaUDFCodegen.CacheKey(ByteBuffer.wrap(bytes), specsSeq) + val key = CometScalaUDFCodegen.CacheKey(planId, ByteBuffer.wrap(digest), specsSeq) // Cache lookup and `process` run under one lock to serialize concurrent Tokio callers that // would otherwise race on the kernel's per-batch instance fields. this.synchronized { - val entry = lookupOrCompile(key, bytes, specsSeq) + val entry = lookupOrCompile(key, inputs(1), specsSeq, partitionIndex) val out = CometBatchKernelCodegen.allocateOutput( entry.outputField, @@ -143,14 +172,20 @@ class CometScalaUDFCodegen extends CometUDF with Logging { private def lookupOrCompile( key: CometScalaUDFCodegen.CacheKey, - bytes: Array[Byte], - specs: IndexedSeq[ArrowColumnSpec]): CometScalaUDFCodegen.CacheEntry = { + exprVec: ValueVector, + specs: IndexedSeq[ArrowColumnSpec], + partitionIndex: Int): CometScalaUDFCodegen.CacheEntry = { assert(Thread.holdsLock(this), "lookupOrCompile must run under this.synchronized") - kernelCache.get(key) match { + // A deterministic kernel never reads the partition index, so one instance under the planless + // key serves every plan in the task. Only a kernel with a nondeterministic node is stored per + // plan. + val sharedKey = key.copy(planId = CometScalaUDFCodegen.NoPlan) + kernelCache.get(sharedKey).orElse(kernelCache.get(key)) match { case Some(entry) => CometScalaUDFCodegen.cacheHitCount.incrementAndGet() entry case None => + val bytes = binaryScalar(exprVec, "serialized expression at arg 1") val loader = Option(Thread.currentThread().getContextClassLoader) .getOrElse(classOf[Expression].getClassLoader) val boundExpr = @@ -168,20 +203,32 @@ class CometScalaUDFCodegen extends CometUDF with Logging { } val compiled = CometBatchKernelCodegen.compile(boundExpr, specs) val kernel = compiled.newInstance() - kernel.init(CometScalaUDFCodegen.currentPartitionIndex()) + kernel.init(partitionIndex) val outputField = CometBatchKernelCodegen.toFfiArrowField( "codegen_result", boundExpr.dataType, boundExpr.nullable) val entry = CometScalaUDFCodegen.CacheEntry(compiled, kernel, boundExpr.dataType, outputField) - kernelCache.put(key, entry) + // Walks the tree, because `deterministic` is not transitive everywhere. `Invoke` skips its + // `targetObject`, which is where Spark 4's `make_valid_utf8` puts its input. + val perPlan = boundExpr.exists(!_.deterministic) + kernelCache.put(if (perPlan) key else sharedKey, entry) CometScalaUDFCodegen.compileCount.incrementAndGet() CometScalaUDFCodegen.recordCompiledSignature(specs, boundExpr.dataType) entry } } + /** The value of a binary scalar argument, which arrives as a length-1 vector. */ + private def binaryScalar(v: ValueVector, what: String): Array[Byte] = { + val vec = v.asInstanceOf[VarBinaryVector] + require( + vec.getValueCount >= 1 && !vec.isNull(0), + s"CometScalaUDFCodegen requires a non-null $what") + vec.get(0) + } + /** * Build the compile-time spec for one input Arrow vector. Recurses on complex types. * @@ -192,9 +239,10 @@ class CometScalaUDFCodegen extends CometUDF with Logging { * * Top-level `nullable=true` is hardcoded: the cache key does not specialize on per-batch null * density. Schema-declared nullability still reaches the kernel via `BoundReference.nullable` - * embedded in `bytesKey`, so `BoundReference.doGenCode` elides its own `isNullAt` probe on - * non-null columns. `StructFieldSpec.nullable` reads `field.isNullable` from Arrow metadata, - * which is a schema property and therefore stable across batches. + * embedded in the serialized expression, which the key's digest covers, so + * `BoundReference.doGenCode` elides its own `isNullAt` probe on non-null columns. + * `StructFieldSpec.nullable` reads `field.isNullable` from Arrow metadata, which is a schema + * property and therefore stable across batches. */ private def specFor(v: ValueVector): ArrowColumnSpec = v match { case map: MapVector => @@ -300,21 +348,28 @@ object CometScalaUDFCodegen { } /** - * Partition index for the kernel's `init`. Expressions whose `doGenCode` calls - * `addPartitionInitializationStatement` (`Rand`, `Randn`, `Uuid`) reseed mutable state from - * this. Falls back to 0 when the dispatcher is exercised outside a Spark task (unit tests). + * Plan id for kernels shared by every plan in a task, and for callers that bypass the bridge. */ - private def currentPartitionIndex(): Int = - Option(TaskContext.get()).map(_.partitionId()).getOrElse(0) + private[comet] val NoPlan = -1L + + /** Length in bytes of a [[digest]]. */ + val DigestLength = 32 /** - * Cache key: serialized expression bytes plus per-column compile-time invariants. `hashCode` - * walks `bytesKey` per lookup, so for large ScalaUDF closures it scales with closure size. - * - * TODO(perf-cache-key): if hot, options are a driver-precomputed hash piggybacked through the - * proto, per-instance last-key memoization, or a two-tier cache keyed on the generated source. + * Identifies a closure-serialized bound expression in the kernel cache. The serde computes it + * once on the driver and ships it as arg 0, next to the bytes, so the per-batch lookup hashes + * 32 bytes instead of the whole serialized expression. SHA-256 rather than a cheaper hash, + * because a hit is trusted without comparing the bytes: two expressions whose digests collided + * would share a kernel. + */ + def digest(serialized: Array[Byte]): Array[Byte] = + MessageDigest.getInstance("SHA-256").digest(serialized) + + /** + * Cache key: calling native plan (`NoPlan` for a deterministic expression), the bound + * expression's [[digest]] plus per-column compile-time invariants. */ - final case class CacheKey(bytesKey: ByteBuffer, specs: IndexedSeq[ArrowColumnSpec]) + final case class CacheKey(planId: Long, digest: ByteBuffer, specs: IndexedSeq[ArrowColumnSpec]) /** Snapshot of dispatcher cache counters and current size. */ final case class DispatcherStats(compileCount: Long, cacheHitCount: Long, cacheSize: Int) { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala index c944d3af8cd..e058a5c0916 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala @@ -27,7 +27,7 @@ import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, SortOrder} import org.apache.spark.sql.catalyst.plans.QueryPlan import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, UnknownPartitioning} -import org.apache.spark.sql.execution.SQLExecution +import org.apache.spark.sql.execution.{SparkPlan, SQLExecution} import org.apache.spark.sql.execution.datasources.v2.BatchScanExec import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.vectorized.ColumnarBatch @@ -65,6 +65,14 @@ case class CometIcebergNativeScanExec( @transient nativeIcebergScanMetadata: CometIcebergNativeScanMetadata) extends CometLeafExec { + override def sparkFallback(newChildren: Seq[SparkPlan]): SparkPlan = { + // Native execution rebuilds originalPlan from the live runtimeFilters during partition + // serialization. Reversion skips that path, so carry the executable DPP filters across here. + val restoredScan = originalPlan.copy(runtimeFilters = runtimeFilters) + originalPlan.logicalLink.foreach(restoredScan.setLogicalLink) + restoredScan + } + override val supportsColumnar: Boolean = true override val nodeName: String = "CometIcebergNativeScan" diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala index 877eb84c1f1..83149ced77c 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala @@ -24,14 +24,13 @@ import java.util.concurrent.atomic.AtomicReference import org.apache.spark.TaskContext import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference} +import org.apache.spark.sql.catalyst.expressions.Attribute import org.apache.spark.sql.catalyst.expressions.UnsafeProjection import org.apache.spark.sql.comet.execution.arrow.CometArrowStream import org.apache.spark.sql.comet.util.{Utils => CometUtils} import org.apache.spark.sql.connector.write.{BatchWrite, WriterCommitMessage} import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} -import org.apache.spark.sql.types.BinaryType import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.spark.util.TaskFailureListener @@ -54,8 +53,12 @@ import org.apache.comet.serde.OperatorOuterClass.Operator * @param nativeOp * Template operator carrying the `IcebergWrite` proto. Per-task `partition_id` / * `task_attempt_id` get stamped on a copy at execution time. + * @param originalPlan + * The JVM Iceberg write operator restored when Comet reverts a transition-heavy stage. * @param child * Comet native child (must be a [[CometNativeExec]] so columnar batches flow through FFI). + * @param output + * Commit-message output attributes copied from the JVM Iceberg write plan. * @param batchWrite * Shared with the outer [[IcebergCommitExec]] -- the same instance the strategy materialised * via `write.toBatch`. Used here only to provide the `dataLocation` / partition spec needed by @@ -66,20 +69,15 @@ import org.apache.comet.serde.OperatorOuterClass.Operator */ case class CometIcebergWriteExec( nativeOp: Operator, + @transient override val originalPlan: IcebergWriteExec, child: SparkPlan, + override val output: Seq[Attribute], @transient batchWrite: BatchWrite, @transient table: AnyRef, partitionSpecId: Int) extends CometNativeExec with UnaryExecNode { - override def originalPlan: SparkPlan = child - - // Same output schema as IcebergWriteExec so the outer IcebergCommitExec consumes the - // commit messages identically regardless of which inner exec emitted them. - override def output: Seq[Attribute] = Seq( - AttributeReference(IcebergWriteExec.CommitMessageColumn, BinaryType, nullable = false)()) - // Native exec emits a single Binary column; the surrounding command framework expects rows, so // the outer commit exec calls executeCollect on us. supportsColumnar = false keeps Spark from // inserting a ColumnarToRow that would clash with our (Nil-output-like) row contract. diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometLocalTopKExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometLocalTopKExec.scala index ec3c85191a3..16ecb208819 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometLocalTopKExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometLocalTopKExec.scala @@ -91,6 +91,15 @@ case class CometLocalTopKExec( override val serializedPlanOpt: SerializedPlan) extends CometUnaryExec { + override def sparkFallback(newChildren: Seq[SparkPlan]): SparkPlan = newChildren match { + // Local candidate selection was inserted by Comet. Only the outer TopK owns + // the original Spark operator's offset and projection. + case Seq(restoredChild) => restoredChild + case _ => + throw new CometExec.InvalidSparkFallbackException( + s"CometLocalTopKExec expected one restored child but received ${newChildren.size}") + } + override def outputPartitioning: Partitioning = child.outputPartitioning override def outputOrdering: Seq[SortOrder] = sortOrder diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala index 8ea48c1002d..37694992211 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala @@ -75,6 +75,15 @@ case class CometNativeScanExec( with ShimStreamSourceAwareSparkPlan with CometScanWithPlanData { + override def sparkFallback(newChildren: Seq[SparkPlan]): SparkPlan = { + // AQE rewrites DPP placeholders in the live scan before transition reversion. + // The frozen originalPlan can still contain unexecutable SubqueryAdaptiveBroadcastExecs. + val restoredScan = + originalPlan.copy(partitionFilters = partitionFilters, dataFilters = dataFilters) + originalPlan.logicalLink.foreach(restoredScan.setLogicalLink) + restoredScan + } + override lazy val metadata: Map[String, String] = if (originalPlan != null) originalPlan.metadata else Map.empty diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala index ced525c1869..ec903fc8937 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeWriteExec.scala @@ -19,80 +19,74 @@ package org.apache.spark.sql.comet -import scala.jdk.CollectionConverters._ +import java.util.Date +import org.apache.hadoop.conf.Configuration import org.apache.hadoop.fs.Path import org.apache.hadoop.mapreduce.{Job, TaskAttemptContext, TaskAttemptID, TaskID, TaskType} import org.apache.hadoop.mapreduce.task.TaskAttemptContextImpl +import org.apache.parquet.hadoop.codec.CodecConfig import org.apache.spark.TaskContext -import org.apache.spark.internal.io.{FileCommitProtocol, FileNameSpec} +import org.apache.spark.internal.io.{FileCommitProtocol, FileNameSpec, SparkHadoopWriterUtils} +import org.apache.spark.internal.io.FileCommitProtocol.TaskCommitMessage import org.apache.spark.rdd.RDD import org.apache.spark.sql.SaveMode import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.Attribute import org.apache.spark.sql.comet.execution.arrow.CometArrowStream import org.apache.spark.sql.comet.util.{Utils => CometUtils} import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} +import org.apache.spark.sql.execution.command.{CommandUtils, DataWritingCommandExec} +import org.apache.spark.sql.execution.datasources.{InsertIntoHadoopFsRelationCommand, OutputWriterFactory, WriteFilesExec} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.util.SchemaUtils import org.apache.spark.sql.vectorized.ColumnarBatch -import org.apache.spark.util.Utils +import org.apache.spark.util.{SerializableConfiguration, Utils} -import com.google.protobuf.CodedOutputStream - -import org.apache.comet.CometExecIterator import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.operator.NativeWriteUtils /** - * Comet physical operator for native Parquet write operations with FileCommitProtocol support. - * - * This operator writes data to Parquet files using the native Comet engine. It integrates with - * Spark's FileCommitProtocol to provide atomic writes with proper staging and commit semantics. + * Native Parquet writes on Spark 3.x, where Comet replaces the entire write command. * - * The implementation includes support for Spark's file commit protocol through work_dir, job_id, - * and task_attempt_id parameters that can be set in the operator. When work_dir is set, files are - * written to a temporary location that can be atomically committed later. + * Both execution entry points run Spark's job and task commit lifecycle. Native code writes the + * exact filename supplied by the configured commit protocol, using the prepared Hadoop job. Spark + * 4.0+ uses CometWriteFilesExec and leaves the surrounding command with Spark instead. * - * @param nativeOp - * The native operator representing the write operation (template, will be modified per task) - * @param child - * The child operator providing the data to write - * @param outputPath - * The path where the Parquet file will be written - * @param mode - * The Spark SaveMode governing target-exists behavior (Append / Overwrite / ErrorIfExists / - * Ignore). Comet takes over Spark's DataWritingCommandExec so we must apply these semantics - * here - a direct port of Spark's InsertIntoHadoopFsRelationCommand.run() doInsertion logic. - * @param committer - * FileCommitProtocol for atomic writes. If None, files are written directly. - * @param jobTrackerID - * Unique identifier for this write job + * @param originalPlan + * The JVM data-writing command restored when Comet reverts a transition-heavy stage */ case class CometNativeWriteExec( nativeOp: Operator, + @transient override val originalPlan: DataWritingCommandExec, child: SparkPlan, outputPath: String, mode: SaveMode, - committer: Option[FileCommitProtocol] = None, - jobTrackerID: String = Utils.createTempDir().getName) + committer: FileCommitProtocol, + serializableHadoopConf: SerializableConfiguration, + outputWriterFactory: OutputWriterFactory, + jobTrackerID: String = SparkHadoopWriterUtils.createJobTrackerID(new Date())) extends CometNativeExec with UnaryExecNode { - override def originalPlan: SparkPlan = child + override def output: Seq[Attribute] = child.output - // Accumulator to collect TaskCommitMessages from all tasks - // Must be eagerly initialized on driver, not lazy - @transient private val taskCommitMessagesAccum = - sparkContext.collectionAccumulator[FileCommitProtocol.TaskCommitMessage]("taskCommitMessages") - - override def serializedPlanOpt: SerializedPlan = { - val size = nativeOp.getSerializedSize - val bytes = new Array[Byte](size) - val codedOutput = CodedOutputStream.newInstance(bytes) - nativeOp.writeTo(codedOutput) - codedOutput.checkNoSpaceLeft() - SerializedPlan(Some(bytes)) + override def sparkFallback(newChildren: Seq[SparkPlan]): SparkPlan = newChildren match { + case Seq(newInput) => + val restoredCommandChild = originalPlan.child match { + case writeFiles: WriteFilesExec => writeFiles.withNewChildren(Seq(newInput)) + case _ => newInput + } + originalPlan.withNewChildren(Seq(restoredCommandChild)) + case _ => + throw new CometExec.InvalidSparkFallbackException( + s"${getClass.getSimpleName} expected one reverted input but received ${newChildren.size}") } + override def serializedPlanOpt: SerializedPlan = + SerializedPlan(Some(CometExec.serializeNativePlan(nativeOp))) + override def withNewChildInternal(newChild: SparkPlan): SparkPlan = copy(child = newChild) @@ -104,219 +98,163 @@ case class CometNativeWriteExec( "rows_written" -> SQLMetrics.createMetric(sparkContext, "number of written rows")) override def doExecute(): RDD[InternalRow] = { - // Setup job if committer is present - committer.foreach { c => - val jobContext = createJobContext() - c.setupJob(jobContext) - } - - // Execute the native write with commit protocol - val resultRDD = doExecuteColumnar() + executeWriteAndCommit() + sparkContext.emptyRDD[InternalRow] + } - // Force execution by consuming all batches - resultRDD - .mapPartitions { iter => - iter.foreach(_.close()) - Iterator.empty - } - .count() + override def doExecuteColumnar(): RDD[ColumnarBatch] = { + executeWriteAndCommit() + sparkContext.emptyRDD[ColumnarBatch] + } - // Extract write statistics from metrics - val filesWritten = metrics("files_written").value - val bytesWritten = metrics("bytes_written").value - val rowsWritten = metrics("rows_written").value + /** The command this exec replaces, which holds the table this write updates, if any. */ + private def command: InsertIntoHadoopFsRelationCommand = originalPlan.cmd match { + case cmd: InsertIntoHadoopFsRelationCommand => cmd + case other => + throw new IllegalStateException( + s"${getClass.getSimpleName} replaced an unexpected write command: ${other.getClass}") + } - // Collect TaskCommitMessages from accumulator - val commitMessages = taskCommitMessagesAccum.value.asScala.toSeq + private def executeWriteAndCommit(): Unit = { + val cmd = command + // Like InsertIntoHadoopFsRelationCommand.run, reject duplicate columns before anything else. + SchemaUtils.checkColumnNameDuplication( + cmd.outputColumnNames, + session.sessionState.conf.caseSensitiveAnalysis) - // Commit job with collected TaskCommitMessages - committer.foreach { c => - val jobContext = createJobContext() - try { - c.commitJob(jobContext, commitMessages) - logInfo( - s"Successfully committed write job to $outputPath: " + - s"$filesWritten files, $bytesWritten bytes, $rowsWritten rows") - } catch { - case e: Exception => - logError("Failed to commit job, aborting", e) - c.abortJob(jobContext) - throw e - } + if (!prepareOutputPathForMode()) { + logInfo(s"Skipping insertion into $outputPath - already exists (SaveMode.$mode)") + return } - // Return empty RDD as write operations don't return data - sparkContext.emptyRDD[InternalRow] + val job = Job.getInstance(new Configuration(serializableHadoopConf.value)) + // Like FileFormatWriter, only abort after setupJob has succeeded. + committer.setupJob(job) + Utils.tryWithSafeFinallyAndFailureCallbacks(block = { + // Include configuration changes made by setupJob in the task contexts. + val commitMessages = runNativeWriteJob(new SerializableConfiguration(job.getConfiguration)) + committer.commitJob(job, commitMessages.toSeq) + logInfo( + s"Successfully committed native write job to $outputPath: " + + s"${metrics("files_written").value} files, " + + s"${metrics("bytes_written").value} bytes, ${metrics("rows_written").value} rows") + })(catchBlock = committer.abortJob(job)) + // Outside the block above: the job has committed, so a failure here must not abort it. + refreshAfterWrite(cmd) } - override def doExecuteColumnar(): RDD[ColumnarBatch] = { - // Comet replaces DataWritingCommandExec entirely, so Spark's - // InsertIntoHadoopFsRelationCommand.run() never runs. That method is where Spark handles - // SaveMode semantics (path-exists check, delete-before-Overwrite, Ignore short-circuit) - - // port the non-partitioned, non-catalog branch of that logic here. See Spark 3.5's - // InsertIntoHadoopFsRelationCommand.run doInsertion match. This runs on the driver before - // any executor tasks fire, mirroring where Spark does the delete. - if (!prepareOutputPathForMode()) { - logInfo(s"Skipping insertion into $outputPath - already exists (SaveMode.$mode)") - return sparkContext.emptyRDD[ColumnarBatch] - } + /** + * Repeats what InsertIntoHadoopFsRelationCommand.run does once the write has committed, since + * this exec replaces that command. Partition metadata updates are left out because partitioned + * writes never reach this exec. + */ + private def refreshAfterWrite(cmd: InsertIntoHadoopFsRelationCommand): Unit = { + val fs = cmd.outputPath.getFileSystem(serializableHadoopConf.value) + // Refresh the files the table's relation has listed. + cmd.fileIndex.foreach(_.refresh()) + // Recache any cached plan that reads the output path. + session.sharedState.cacheManager.recacheByPath(session, cmd.outputPath, fs) + // Update or clear the table stats, depending on spark.sql.statistics.size.autoUpdate. + cmd.catalogTable.foreach(CommandUtils.updateTableStats(session, _)) + } - // Get the input data from the child operator - val childRDD = if (child.supportsColumnar) { + private def runNativeWriteJob( + hadoopConf: SerializableConfiguration): Array[TaskCommitMessage] = { + val childRDD: RDD[ColumnarBatch] = if (child.supportsColumnar) { child.executeColumnar() } else { - // If child doesn't support columnar, convert to columnar child.execute().mapPartitionsInternal { _ => - // TODO this could delegate to CometRowToColumnar, but maybe Comet - // does not need to support this case? throw new UnsupportedOperationException( "Row-based child operators not yet supported for native write") } } - // Capture metadata before the transformation - val numPartitions = childRDD.getNumPartitions - val numOutputCols = child.output.length + // SPARK-23271: a zero-partition input would spawn no task and therefore write no file at all, + // so the output directory would carry no schema for readers. Swap in a single empty + // partition, as Spark's FileFormatWriter and CometWriteFilesExec do. + val writeRDD = if (childRDD.getNumPartitions == 0) { + sparkContext.parallelize(Seq.empty[ColumnarBatch], 1) + } else { + childRDD + } + + // The task closure below must not touch `this`, which holds `nativeOp` plus the whole + // converted child subtree; capture what it needs by value instead, as CometWriteFilesExec + // does. + val numPartitions = writeRDD.getNumPartitions + val childSchema = CometUtils.fromAttributes(child.output) + val capturedNativeOp = nativeOp val capturedCommitter = committer val capturedJobTrackerID = jobTrackerID - val capturedNativeOp = nativeOp - val capturedAccumulator = taskCommitMessagesAccum // Capture accumulator for use in tasks - - // Execute native write operation with task-level commit protocol - childRDD.mapPartitionsInternal { iter => - val partitionId = org.apache.spark.TaskContext.getPartitionId() - val taskAttemptId = org.apache.spark.TaskContext.get().taskAttemptId() - - // Setup task-level commit protocol if provided - val (workDir, taskContext, _) = capturedCommitter - .map { committer => - val taskContext = - createTaskContext(capturedJobTrackerID, partitionId, taskAttemptId.toInt) - - // Setup task - this creates the temporary working directory - committer.setupTask(taskContext) - - // Get the work directory for temp files - // Spark 4.1 made the (taskContext, dir, ext: String) overload throw by default; - // the FileNameSpec overload is the supported one and exists in 3.4+. - val workPath = committer.newTaskTempFile(taskContext, None, FileNameSpec("", "")) - val workDir = new Path(workPath).getParent.toString - - (Some(workDir), Some((committer, taskContext)), null) - } - .getOrElse((None, None, null)) - - // Modify the native operator to include task-specific parameters - val modifiedNativeOp = if (workDir.isDefined) { - val parquetWriter = capturedNativeOp.getParquetWriter.toBuilder - .setWorkDir(workDir.get) - .setJobId(capturedJobTrackerID) - .setTaskAttemptId(taskAttemptId.toInt) - .build() - - capturedNativeOp.toBuilder.setParquetWriter(parquetWriter).build() - } else { - capturedNativeOp - } - - val nativeMetrics = CometMetricNode.fromCometPlan(this) - // Register before CometExecIterator so completion listeners run after iterator close - // (Spark runs task completion callbacks in reverse registration order). - Option(TaskContext.get()).foreach(nativeMetrics.reportNativeWriteOutputMetrics) - - val size = modifiedNativeOp.getSerializedSize - val planBytes = new Array[Byte](size) - val codedOutput = CodedOutputStream.newInstance(planBytes) - modifiedNativeOp.writeTo(codedOutput) - codedOutput.checkNoSpaceLeft() - - val execIterator = new CometExecIterator( - CometExec.newIterId, - CometArrowStream.inputObjects( - iter, - CometUtils.fromAttributes(child.output), - "CometNativeWriteExec"), - numOutputCols, - planBytes, - nativeMetrics, - numPartitions, - partitionId, - None, - Seq.empty) - - // Wrap the iterator to handle task commit/abort and capture TaskCommitMessage - new Iterator[ColumnarBatch] { - private var completed = false - private var thrownException: Option[Throwable] = None - - override def hasNext: Boolean = { - val result = - try { - execIterator.hasNext - } catch { - case e: Throwable => - thrownException = Some(e) - handleTaskEnd() - throw e - } - - if (!result && !completed) { - handleTaskEnd() - } - - result - } - - override def next(): ColumnarBatch = { - try { - execIterator.next() - } catch { - case e: Throwable => - thrownException = Some(e) - handleTaskEnd() - throw e - } - } - - private def handleTaskEnd(): Unit = { - if (!completed) { - completed = true - - // Handle commit or abort based on whether an exception was thrown - taskContext.foreach { case (committer, ctx) => - try { - if (thrownException.isEmpty) { - // Commit the task and add message to accumulator - val message = committer.commitTask(ctx) - capturedAccumulator.add(message) - logDebug(s"Task ${ctx.getTaskAttemptID} committed successfully") - } else { - // Abort the task - committer.abortTask(ctx) - val exMsg = thrownException.get.getMessage - logWarning(s"Task ${ctx.getTaskAttemptID} aborted due to exception: $exMsg") - } - } catch { - case e: Exception => - // Log the commit/abort exception but don't mask the original exception - logError(s"Error during task commit/abort: ${e.getMessage}", e) - if (thrownException.isEmpty) { - // If no original exception, propagate the commit/abort exception - throw e - } + val writerFactory = outputWriterFactory + val nativeMetrics = CometMetricNode.fromCometPlan(this) + val commitMessages = new Array[TaskCommitMessage](numPartitions) + + sparkContext.runJob( + writeRDD, + (context: TaskContext, batches: Iterator[ColumnarBatch]) => { + val taskContext = + CometNativeWriteExec.createTaskContext(hadoopConf.value, capturedJobTrackerID, context) + capturedCommitter.setupTask(taskContext) + // Guard filename allocation, native iterator construction, execution, cleanup and commit. + // Spark's helper preserves the original error if abortTask also fails. + Utils.tryWithSafeFinallyAndFailureCallbacks(block = { + // Mirrors FileFormatWriter's EmptyDirectoryDataWriter case: an empty input still writes + // one file from partition 0 so that the output carries the schema, but every other + // empty partition produces no file at all. + if (context.partitionId() == 0 || batches.hasNext) { + val extension = writerFactory.getFileExtension(taskContext) + val filePath = capturedCommitter.newTaskTempFile( + taskContext, + None, + FileNameSpec("", "-c000" + extension)) + NativeWriteUtils.checkNativeWriteDestination(filePath) + // `extension` names the codec in the job configuration. The writer uses that codec + // too, so the file name and the footer agree. + val codec = CodecConfig.from(taskContext).getCodec + val protoCodec = NativeWriteUtils + .protoCompressionCodec(codec.name()) + .getOrElse(throw new UnsupportedOperationException( + s"Comet's native Parquet writer cannot write $codec")) + val writer = capturedNativeOp.getParquetWriter.toBuilder + .setOutputPath(filePath) + .setCompression(protoCodec) + .build() + val taskOp = capturedNativeOp.toBuilder.setParquetWriter(writer).build() + + // Register before the iterator so this listener runs after its cleanup. + nativeMetrics.reportNativeWriteOutputMetrics(context) + val execIterator = CometExec.getCometIterator( + CometArrowStream.inputObjects(batches, childSchema, "CometNativeWriteExec"), + childSchema.length, + taskOp, + nativeMetrics, + numPartitions, + context.partitionId(), + None, + Seq.empty) + + // Close before committing. A failed write must remain the primary error even if + // native teardown (including the final metrics update) also throws. + Utils.tryWithSafeFinally { + while (execIterator.hasNext) { + execIterator.next().close() } + } { + execIterator.close() } } - } - } - } - } - - /** Create a JobContext for the write job */ - private def createJobContext(): Job = { - val job = Job.getInstance() - job.setJobID(new org.apache.hadoop.mapreduce.JobID(jobTrackerID, 0)) - job + // `hasNext` already ran an empty partition's child to completion, so there is nothing + // to drain or release when no file was written. + capturedCommitter.commitTask(taskContext) + })(catchBlock = capturedCommitter.abortTask(taskContext)) + }, + writeRDD.partitions.indices, + (index, message: TaskCommitMessage) => { + committer.onTaskCommit(message) + commitMessages(index) = message + }) + commitMessages } /** @@ -329,7 +267,7 @@ case class CometNativeWriteExec( */ private def prepareOutputPathForMode(): Boolean = { val path = new Path(outputPath) - val hadoopConf = sparkContext.hadoopConfiguration + val hadoopConf = serializableHadoopConf.value val fs = path.getFileSystem(hadoopConf) val qualifiedOutputPath = path.makeQualified(fs.getUri, fs.getWorkingDirectory) @@ -343,11 +281,7 @@ case class CometNativeWriteExec( true case SaveMode.Overwrite => if (fs.exists(qualifiedOutputPath)) { - val deleted = committer match { - case Some(c) => c.deleteWithJob(fs, qualifiedOutputPath, true) - case None => fs.delete(qualifiedOutputPath, true) - } - if (!deleted) { + if (!committer.deleteWithJob(fs, qualifiedOutputPath, true)) { throw QueryExecutionErrors.cannotClearOutputDirectoryError(qualifiedOutputPath) } } @@ -357,15 +291,27 @@ case class CometNativeWriteExec( } } - /** Create a TaskAttemptContext for a specific task */ +} + +object CometNativeWriteExec { + + /** + * Match FileFormatWriter's Hadoop task identifiers and configuration. Kept off the exec node so + * that the task closure does not capture it. + */ private def createTaskContext( - jobId: String, - partitionId: Int, - attemptNumber: Int): TaskAttemptContext = { - val job = Job.getInstance() - val taskAttemptID = new TaskAttemptID( - new TaskID(new org.apache.hadoop.mapreduce.JobID(jobId, 0), TaskType.REDUCE, partitionId), - attemptNumber) - new TaskAttemptContextImpl(job.getConfiguration, taskAttemptID) + conf: Configuration, + jobTrackerID: String, + context: TaskContext): TaskAttemptContext = { + val hadoopConf = new Configuration(conf) + val jobId = SparkHadoopWriterUtils.createJobID(jobTrackerID, context.stageId()) + val taskId = new TaskID(jobId, TaskType.MAP, context.partitionId()) + val attemptId = new TaskAttemptID(taskId, context.taskAttemptId().toInt & Integer.MAX_VALUE) + hadoopConf.set("mapreduce.job.id", jobId.toString) + hadoopConf.set("mapreduce.task.id", taskId.toString) + hadoopConf.set("mapreduce.task.attempt.id", attemptId.toString) + hadoopConf.setBoolean("mapreduce.task.ismap", true) + hadoopConf.setInt("mapreduce.task.partition", 0) + new TaskAttemptContextImpl(hadoopConf, attemptId) } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala index ee3cbd61c63..e4c24583337 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWindowGroupLimitExec.scala @@ -102,8 +102,8 @@ object CometWindowGroupLimitExec extends CometOperatorSerde[SparkPlan] { // The same byte equality decides RANK and DENSE_RANK ties, but it runs on keys that the native // planner has normalized: a `FLOAT` or `DOUBLE` order key, scalar or nested in an array or // struct, has its NaNs folded together and its signed zeros tied first, as Spark does, so no - // float key needs a guard here. Strict floating-point mode declines the nested keys that can - // hold a null element or field, through `CometSortOrder`. + // float key needs a guard here. `CometSortOrder` declines `ASC NULLS LAST` and + // `DESC NULLS FIRST` on a nested key that can hold a null element or field, in every mode. val childOutput = op.children.head.output val partitionProtos = fields.partitionSpec.map(e => e -> exprToProto(e, childOutput)) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala index 305593e609d..7721cb7c2a5 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometWriteFilesExec.scala @@ -129,10 +129,10 @@ case class CometWriteFilesExec( // SPARK-23271: a zero-partition input would spawn no task and therefore write no file at all, // so the output directory would carry no schema for readers. Spark's own WriteFilesExec swaps - // in a dummy single-partition RDD for exactly this case. AQE reaches it by collapsing a - // completed empty shuffle into a CometEmptyRelationExec, which is why - // CometDataWritingCommand declines empty-relation inputs on the Spark 3.x path (#5303): its - // writer has nowhere to put this swap. See CometEmptyRelationParquetWriterSuite. + // in a dummy single-partition RDD for exactly this case, and CometNativeWriteExec does the + // same on the Spark 3.x path (#5303). AQE reaches it by collapsing a completed empty shuffle + // into a CometEmptyRelationExec, so a write over a native empty relation stays native and + // still leaves a schema-only file. See CometEmptyRelationParquetWriterSuite. val writeRDD = if (childRDD.getNumPartitions == 0) { sparkContext.parallelize(Seq.empty[ColumnarBatch], 1) } else { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/DecimalPrecision.scala b/spark/src/main/scala/org/apache/spark/sql/comet/DecimalPrecision.scala index 08efcfc9c9c..14ab3554776 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/DecimalPrecision.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/DecimalPrecision.scala @@ -40,6 +40,27 @@ import org.apache.spark.sql.types.DecimalType */ object DecimalPrecision { def promote(expr: Expression): Expression = { + // `transformUp` walks the whole tree and calls `mapChildren` at every node, which allocates + // through `children.map` / `mapProductIterator` even when no case matches (the specialized + // `mapChildren` overrides do skip the node copy via `fastEquals`, so the cost is the traversal + // and its allocations, not a rebuild). The serde calls this once per expression it converts, + // so skip it when there is nothing to rewrite. + if (containsDecimalArithmetic(expr)) rewrite(expr) else expr + } + + /** + * Whether [[promote]] could rewrite anything in this tree. Deliberately a superset of the rule: + * it ignores the operands, so any arithmetic node the rule handles that produces a decimal + * qualifies. Wrapping a child in `CheckOverflow` does not change that child's data type, so a + * tree without such a node cannot grow one part way through the rewrite. + */ + private def containsDecimalArithmetic(expr: Expression): Boolean = expr.exists { + case e @ (_: Add | _: Subtract | _: Multiply | _: Divide | _: Remainder) => + e.dataType.isInstanceOf[DecimalType] + case _ => false + } + + private def rewrite(expr: Expression): Expression = { expr.transformUp { // Recursive exprToProto calls can re-promote every decimal binary operator below. Collapse // only equivalent wrappers so the shared promotion rule remains idempotent. diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/IcebergCommitExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/IcebergCommitExec.scala index a2323e0acce..e066cd30ab3 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/IcebergCommitExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/IcebergCommitExec.scala @@ -28,7 +28,7 @@ import org.apache.spark.sql.execution.{SparkPlan, SQLExecution, UnaryExecNode} import org.apache.spark.sql.execution.datasources.v2.V2CommandExec import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} -import org.apache.comet.iceberg.{IcebergDriverMetricsShim, IcebergReflection} +import org.apache.comet.iceberg.{DeltaCommand, IcebergDriverMetricsShim, IcebergReflection} /** * Driver-side committer for Comet's split-operator Iceberg V2 write. @@ -38,7 +38,8 @@ case class IcebergCommitExec( @transient batchWrite: BatchWrite, @transient write: Write, @transient refreshCache: IcebergCommitExec.RefreshCache, - child: SparkPlan) + child: SparkPlan, + command: Option[DeltaCommand] = None) extends V2CommandExec with UnaryExecNode with Logging { @@ -95,9 +96,9 @@ case class IcebergCommitExec( cause) // The BatchWrite contract expects a job-level abort. Iceberg's own `SparkWrite.abort` // only deletes files after a cleanable *commit* failure and skips cleanup otherwise, so - // the data files of tasks that completed before the job failed would stay behind even - // though nothing can reference them (no commit was attempted). Delete them here; the - // failed task's own files are cleaned up by the task itself. + // the files of tasks that completed before the job failed would stay behind even though + // nothing can reference them (no commit was attempted). Delete newly written data/delete + // files here; the failed task's own files are cleaned up by the task itself. val failure = abortAfter(completed, cause) try deleteCompletedTaskFiles(completed) catch { @@ -110,7 +111,7 @@ case class IcebergCommitExec( try { messages.foreach(batchWrite.onDataWriterCommit) - IcebergWriteSummaryShim.commit(batchWrite, messages, child) + IcebergWriteSummaryShim.commit(batchWrite, messages, child, command) logInfo(s"Iceberg commit succeeded with ${messages.length} task message(s)") } catch { case cause: Throwable => @@ -143,7 +144,7 @@ case class IcebergCommitExec( if (locations.nonEmpty) { val io = IcebergReflection .getOuterSparkWrite(batchWrite) - .flatMap(IcebergReflection.getTableFromSparkWrite) + .flatMap(IcebergReflection.getTableFromIcebergWrite) .flatMap(IcebergReflection.getTableIO) io match { case Some(fileIO) => @@ -153,7 +154,7 @@ case class IcebergCommitExec( s"job abort, ${completed.length} completed task(s)") case None => logWarning( - s"Could not resolve the table FileIO; leaving ${locations.size} data file(s) from " + + s"Could not resolve the table FileIO; leaving ${locations.size} written file(s) from " + "completed tasks of a failed write job for remove_orphan_files") } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/IcebergWriteExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/IcebergWriteExec.scala index 90bfdaab048..e18ff4ac5dd 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/IcebergWriteExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/IcebergWriteExec.scala @@ -31,7 +31,7 @@ import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.types.{BinaryType, StructField, StructType} import org.apache.spark.util.Utils -import org.apache.comet.iceberg.ReplaceDataDispatchInfo +import org.apache.comet.iceberg.{IcebergDeltaWriterShim, IcebergSemanticMetricsShim, IcebergWriteDispatch, PlainIcebergWrite, PositionDeltaWrite, ReplaceDataDispatchInfo, ReplaceDataWrite} /** * Executor-side file writer for Comet's split-operator Iceberg V2 write. @@ -41,7 +41,7 @@ case class IcebergWriteExec( @transient batchWrite: BatchWrite, override val output: Seq[Attribute], child: SparkPlan, - replaceDataDispatch: Option[ReplaceDataDispatchInfo] = None) + dispatch: IcebergWriteDispatch = PlainIcebergWrite) extends UnaryExecNode { // Spark already adds a distribution for the V2 write; adding another here is redundant. @@ -51,7 +51,12 @@ case class IcebergWriteExec( // through the driver-side commit report, and its DataWriters expose no task-level values. // Registering them here too would double-count them in the SQL UI. override lazy val metrics: Map[String, SQLMetric] = - Map("numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows")) + Map("numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows")) ++ + (dispatch match { + case PositionDeltaWrite(info) => + IcebergSemanticMetricsShim.deltaMetrics(sparkContext, info.command) + case _ => Map.empty[String, SQLMetric] + }) override protected def doExecute(): RDD[InternalRow] = { val rdd = { @@ -72,7 +77,8 @@ case class IcebergWriteExec( val rowsMetric = longMetric("numOutputRows") val schemaTypes = output.map(_.dataType).toArray - val capturedReplaceDataDispatch = replaceDataDispatch + val capturedDispatch = dispatch + val semanticMetrics = metrics -- Set("numOutputRows") rdd.mapPartitionsInternal { iter => val partId = TaskContext.getPartitionId() val taskId = TaskContext.get().taskAttemptId() @@ -83,7 +89,8 @@ case class IcebergWriteExec( iter, rowsMetric, projection, - capturedReplaceDataDispatch) + capturedDispatch, + semanticMetrics) } } @@ -106,17 +113,40 @@ object IcebergWriteExec { iter: Iterator[InternalRow], rowsMetric: SQLMetric, projection: UnsafeProjection, - replaceDataDispatch: Option[ReplaceDataDispatchInfo]): Iterator[InternalRow] = { + dispatch: IcebergWriteDispatch, + semanticMetrics: Map[String, SQLMetric] = Map.empty): Iterator[InternalRow] = { val iterWithMetrics = new IteratorWithMetrics(iter, rowsMetric) // Serialization happens inside the guarded block: if the commit message cannot be // serialised the task must abort so the already-finalised data files get deleted. val messageBytes = Utils.tryWithSafeFinallyAndFailureCallbacks(block = { - if (replaceDataDispatch.isDefined) { - runReplaceDataWriter(writer, iterWithMetrics, replaceDataDispatch.get) - } else { - while (iterWithMetrics.hasNext) { - writer.write(iterWithMetrics.next()) - } + dispatch match { + case PlainIcebergWrite => + while (iterWithMetrics.hasNext) writer.write(iterWithMetrics.next()) + case ReplaceDataWrite(info) => + runReplaceDataWriter(writer, iterWithMetrics, info) + case PositionDeltaWrite(info) => + val deltaWriter = writer match { + case delta: org.apache.spark.sql.connector.write.DeltaWriter[_] => + delta.asInstanceOf[org.apache.spark.sql.connector.write.DeltaWriter[InternalRow]] + case other => + throw new IllegalStateException( + s"Expected DeltaWriter for position delta write, got ${other.getClass.getName}") + } + var numUpdatedRows = 0L + var numDeletedRows = 0L + while (iterWithMetrics.hasNext) { + val row = iterWithMetrics.next() + val operation = row.getInt(info.operationOrdinal) + if (operation == info.operationCodes.delete) { + numDeletedRows += 1L + } else if (operation == info.operationCodes.update || + info.operationCodes.reinsert.contains(operation)) { + numUpdatedRows += 1L + } + IcebergDeltaWriterShim.writeOperation(deltaWriter, operation, row, info) + } + semanticMetrics.get("numUpdatedRows").foreach(_.add(numUpdatedRows)) + semanticMetrics.get("numDeletedRows").foreach(_.add(numDeletedRows)) } serializeMessage(writer.commit()) })( diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala index 00c4de38f6d..53b4ae7888d 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala @@ -29,7 +29,7 @@ import org.apache.arrow.vector._ import org.apache.arrow.vector.complex._ import org.apache.arrow.vector.util.OversizedAllocationException import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{SpecializedGetters, UnsafeArrayData, UnsafeRow} +import org.apache.spark.sql.catalyst.expressions.{SpecializedGetters, UnsafeArrayData, UnsafeMapData, UnsafeRow} import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.execution.vectorized.{ConstantColumnVector, OffHeapColumnVector, OnHeapColumnVector, WritableColumnVector} @@ -151,9 +151,17 @@ class ArrowWriter(val root: VectorSchemaRoot, fields: Array[ArrowFieldWriter]) { def write(row: InternalRow): Unit = { var i = 0 - while (i < fields.length) { - fields(i).writeUnsafe(row, i) - i += 1 + row match { + case unsafe: UnsafeRow => + while (i < fields.length) { + fields(i).writeUnsafeRowField(unsafe, i) + i += 1 + } + case _ => + while (i < fields.length) { + fields(i).write(row, i) + i += 1 + } } count += 1 } @@ -212,26 +220,62 @@ private[arrow] object ArrowFieldWriter { // Below this many rows the per-value loop is faster. val MinOnHeapBulkCopyRows = 32 + // Below this many bytes, copying a word at a time beats Unsafe.copyMemory's checks and call. + // Most strings, and the elements of most arrays, are shorter. + private final val MaxWordCopyBytes = 64 + + /** Copies `length` bytes from `srcOffset` in `src` to the native address `dst`. */ + def copyMemory(src: AnyRef, srcOffset: Long, dst: Long, length: Long): Unit = { + if (length > MaxWordCopyBytes || !Platform.unaligned()) { + Platform.copyMemory(src, srcOffset, null, dst, length) + } else { + var i = 0L + while (i + 8 <= length) { + Platform.putLong(null, dst + i, Platform.getLong(src, srcOffset + i)) + i += 8 + } + if (i + 4 <= length) { + Platform.putInt(null, dst + i, Platform.getInt(src, srcOffset + i)) + i += 4 + } + if (i + 2 <= length) { + Platform.putShort(null, dst + i, Platform.getShort(src, srcOffset + i)) + i += 2 + } + if (i < length) { + Platform.putByte(null, dst + i, Platform.getByte(src, srcOffset + i)) + } + } + } + /** Spark's own writable vectors, whose storage layout the columnar fast paths rely on. */ def isSparkVector(input: ColumnVector): Boolean = input.isInstanceOf[OnHeapColumnVector] || input.isInstanceOf[OffHeapColumnVector] + // Arrow's setOne calls Unsafe.setMemory, which the JIT does not inline, so shorter runs of whole + // bytes are set one at a time. + private final val MinSetOneBytes = 64 + /** Marks bits `[start, start + numRows)` of `validity` as valid. */ def setValid(validity: ArrowBuf, start: Int, numRows: Int): Unit = { val end = start + numRows - var i = start - while (i < end && (i & 7) != 0) { - BitVectorHelper.setBit(validity, i.toLong) - i += 1 - } - val alignedEnd = end & ~7 - if (i < alignedEnd) { - validity.setOne((i >> 3).toLong, ((alignedEnd - i) >> 3).toLong) - i = alignedEnd - } - while (i < end) { - BitVectorHelper.setBit(validity, i.toLong) - i += 1 + val firstByte = (start + 7) >> 3 + val endByte = end >> 3 + if (firstByte >= endByte) { + // No whole byte, as for most of a map's entries. + writeBits(validity, start, -1L, numRows) + } else { + writeBits(validity, start, -1L, (firstByte << 3) - start) + if (endByte - firstByte >= MinSetOneBytes) { + validity.setOne(firstByte.toLong, (endByte - firstByte).toLong) + } else { + var b = firstByte + while (b < endByte) { + validity.setByte(b.toLong, 0xff) + b += 1 + } + } + writeBits(validity, endByte << 3, -1L, end - (endByte << 3)) } } @@ -282,6 +326,37 @@ private[arrow] object ArrowFieldWriter { } } + /** Writes the low `numBits` bits of `bits`, at most 64, to `[start, start + numBits)`. */ + def writeBits(buffer: ArrowBuf, start: Int, bits: Long, numBits: Int): Unit = { + val end = start + numBits + var remaining = bits + var i = start + while (i < end) { + val shift = i & 7 + val n = Math.min(8 - shift, end - i) + val mask = ((1 << n) - 1) << shift + val index = (i >> 3).toLong + buffer.setByte(index, (buffer.getByte(index) & ~mask) | ((remaining.toInt << shift) & mask)) + remaining >>>= n + i += n + } + } + + /** + * Sets the validity of the elements of `array` from bit `start` of `validity`. Spark keeps an + * unsafe array's null bits in 64-bit words after its element count, a set bit marking a null. + */ + def writeArrayValidity(validity: ArrowBuf, start: Int, array: UnsafeArrayData): Unit = { + val numElements = array.numElements() + val nulls = array.getBaseOffset + 8 + var i = 0 + while (i < numElements) { + val nullBits = Platform.getLong(array.getBaseObject, nulls + (i >> 3)) + writeBits(validity, start + i, ~nullBits, Math.min(64, numElements - i)) + i += 64 + } + } + /** * Clears each bit of `validity` in `[start, start + numRows)` whose bit in `parent` is clear. * Whole bytes are combined: the bits before `start` in the first byte belong to rows already @@ -386,12 +461,7 @@ private[arrow] object ArrowFieldWriter { if (numRows == 0) { return } - if (vector.getLastSet < outStart - 1) { - vector.fillEmpties(outStart) - } - while (vector.getValueCapacity < outStart + numRows) { - vector.reallocValidityAndOffsetBuffers() - } + reserveValues(vector, outStart, numRows) val offsets = vector.getOffsetBuffer val dataStart = vector.getStartOffset(outStart) val hasNull = input.hasNull @@ -413,10 +483,7 @@ private[arrow] object ArrowFieldWriter { } next = offset + length end += length - if (end > Integer.MAX_VALUE) { - throw new OversizedAllocationException( - s"Arrow variable-width data would exceed ${Integer.MAX_VALUE} bytes") - } + checkDataEnd(end) } } offsets.setInt((outStart + i + 1).toLong * BaseVariableWidthVector.OFFSET_WIDTH, end.toInt) @@ -467,12 +534,7 @@ private[arrow] object ArrowFieldWriter { if (numRows == 0) { return } - if (vector.getLastSet < outStart - 1) { - vector.fillEmpties(outStart) - } - while (vector.getValueCapacity < outStart + numRows) { - vector.reallocValidityAndOffsetBuffers() - } + reserveValues(vector, outStart, numRows) val offsets = vector.getOffsetBuffer val ids = input.getDictionaryIds val hasNull = input.hasNull @@ -523,12 +585,30 @@ private[arrow] object ArrowFieldWriter { vector.setLastSet(outStart + numRows - 1) } - /** The data buffer of `vector`, grown to hold at least `end` bytes. */ - private def reserveData(vector: BaseVariableWidthVector, end: Long): ArrowBuf = { + /** Rejects variable-width data ending at `end`, past what Arrow's 32-bit offsets address. */ + def checkDataEnd(end: Long): Unit = { if (end > Integer.MAX_VALUE) { throw new OversizedAllocationException( s"Arrow variable-width data would exceed ${Integer.MAX_VALUE} bytes") } + } + + /** + * Grows the validity and offset buffers of `vector` to hold values `[outStart, outStart + + * numValues)`, after filling in the offsets of any values skipped since the last one set. + */ + def reserveValues(vector: BaseVariableWidthVector, outStart: Int, numValues: Int): Unit = { + if (vector.getLastSet < outStart - 1) { + vector.fillEmpties(outStart) + } + while (vector.getValueCapacity < outStart + numValues) { + vector.reallocValidityAndOffsetBuffers() + } + } + + /** The data buffer of `vector`, grown to hold at least `end` bytes. */ + def reserveData(vector: BaseVariableWidthVector, end: Long): ArrowBuf = { + checkDataEnd(end) if (vector.getDataBuffer.capacity < end) { vector.reallocDataBuffer(end) } @@ -610,8 +690,20 @@ private[arrow] abstract class ArrowFieldWriter { count += 1 } - def writeUnsafe(input: SpecializedGetters, ordinal: Int): Unit = { - write(input, ordinal) + /** + * Appends field `ordinal` of `row`. Most rows Spark produces are unsafe rows, and writers that + * can read one's memory directly override this. + */ + private[arrow] def writeUnsafeRowField(row: UnsafeRow, ordinal: Int): Unit = write(row, ordinal) + + /** Appends every element of `array`, which writers that can copy them in bulk override. */ + private[arrow] def writeArrayElements(array: UnsafeArrayData): Unit = { + val numElements = array.numElements() + var i = 0 + while (i < numElements) { + write(array, i) + i += 1 + } } def writeCol(input: ColumnarArray): Unit = { @@ -691,16 +783,22 @@ private[arrow] abstract class ArrowFieldWriter { } } -private[arrow] abstract class FixedWidthArrowFieldWriter extends ArrowFieldWriter { +/** + * `vector` is `valueVector`. The methods called per value read it through this field rather than + * the subclass's accessor, a virtual call. + */ +private[arrow] abstract class FixedWidthArrowFieldWriter(vector: BaseFixedWidthVector) + extends ArrowFieldWriter { import ArrowFieldWriter._ override def valueVector: BaseFixedWidthVector protected def setValueUnsafe(input: SpecializedGetters, ordinal: Int): Unit + // Arrow keeps a fixed-width vector's capacity in a field, so checking it per value is cheap. protected def ensureCapacity(inputNumElements: Int): Unit = { - while (valueVector.getValueCapacity < inputNumElements) { - valueVector.reAlloc() + while (vector.getValueCapacity < inputNumElements) { + vector.reAlloc() } } @@ -841,26 +939,75 @@ private[arrow] abstract class FixedWidthArrowFieldWriter extends ArrowFieldWrite } override def setNull(): Unit = { - valueVector.setNull(count) + vector.setNull(count) } protected def setNullUnsafe(): Unit = { - BitVectorHelper.unsetBit(valueVector.getValidityBuffer, count) + BitVectorHelper.unsetBit(vector.getValidityBuffer, count) } - override def writeUnsafe(input: SpecializedGetters, ordinal: Int): Unit = { - if (input.isNullAt(ordinal)) { + // The width of this type's values when Spark's unsafe formats hold the bits Arrow stores, at the + // vector's type width and little-endian: an unsafe row in an 8-byte slot, an unsafe array back + // to back. 0 for a type whose values need converting. + private val unsafeValueWidth: Int = vector match { + case _ if !LittleEndian => 0 + case _: TinyIntVector | _: SmallIntVector | _: IntVector | _: DateDayVector | + _: IntervalYearVector | _: Float4Vector | _: BigIntVector | _: TimeStampMicroTZVector | + _: TimeStampMicroVector | _: DurationVector | _: TimeNanoVector | _: Float8Vector => + vector.getTypeWidth + case _ => 0 + } + + override private[arrow] def writeUnsafeRowField(row: UnsafeRow, ordinal: Int): Unit = { + ensureCapacity(count + 1) + if (row.isNullAt(ordinal)) { setNullUnsafe() + } else if (unsafeValueWidth == 0) { + setValueUnsafe(row, ordinal) } else { - setValueUnsafe(input, ordinal) + val target = vector.getDataBufferAddress + count.toLong * unsafeValueWidth + unsafeValueWidth match { + case 8 => Platform.putLong(null, target, row.getLong(ordinal)) + case 4 => Platform.putInt(null, target, row.getInt(ordinal)) + case 2 => Platform.putShort(null, target, row.getShort(ordinal)) + case _ => Platform.putByte(null, target, row.getByte(ordinal)) + } + BitVectorHelper.setBit(vector.getValidityBuffer, count.toLong) } count += 1 } + override private[arrow] def writeArrayElements(array: UnsafeArrayData): Unit = { + val numElements = array.numElements() + ensureCapacity(count + numElements) + if (unsafeValueWidth == 0) { + var i = 0 + while (i < numElements) { + if (array.isNullAt(i)) { + setNullUnsafe() + } else { + setValueUnsafe(array, i) + } + count += 1 + i += 1 + } + } else { + // The values are copied in one block, so what lands under a null element is whatever Spark + // left there, as with the columnar copies. + copyMemory( + array.getBaseObject, + array.getBaseOffset + UnsafeArrayData.calculateHeaderPortionInBytes(numElements), + vector.getDataBufferAddress + count.toLong * unsafeValueWidth, + numElements.toLong * unsafeValueWidth) + writeArrayValidity(vector.getValidityBuffer, count, array) + count += numElements + } + } + override def writeColumnSlice(input: ColumnVector, startRow: Int, numRows: Int): Unit = { ensureCapacity(count + numRows) if (copyValues(input, startRow, numRows)) { - writeValidity(valueVector.getValidityBuffer, count, input, startRow, numRows) + writeValidity(vector.getValidityBuffer, count, input, startRow, numRows) count += numRows } else { val slice = new ColumnarArray(input, startRow, numRows) @@ -901,7 +1048,7 @@ private[arrow] abstract class FixedWidthArrowFieldWriter extends ArrowFieldWrite } private[arrow] class BooleanWriter(val valueVector: BitVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, if (input.getBoolean(ordinal)) 1 else 0) @@ -949,10 +1096,25 @@ private[arrow] class BooleanWriter(val valueVector: BitVector) ArrowFieldWriter.writeValidity(valueVector.getValidityBuffer, count, input, startRow, numRows) count += numRows } + + // An unsafe array holds a byte per value, each of which becomes a bit. Null elements get a + // clear bit, as in writeColumnSlice. + override private[arrow] def writeArrayElements(array: UnsafeArrayData): Unit = { + val numElements = array.numElements() + ensureCapacity(count + numElements) + val values = valueVector.getDataBuffer + var i = 0 + while (i < numElements) { + ArrowFieldWriter.writeBit(values, count + i, !array.isNullAt(i) && array.getBoolean(i)) + i += 1 + } + ArrowFieldWriter.writeArrayValidity(valueVector.getValidityBuffer, count, array) + count += numElements + } } private[arrow] class ByteWriter(val valueVector: TinyIntVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getByte(ordinal)) @@ -964,7 +1126,7 @@ private[arrow] class ByteWriter(val valueVector: TinyIntVector) } private[arrow] class ShortWriter(val valueVector: SmallIntVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getShort(ordinal)) @@ -976,7 +1138,7 @@ private[arrow] class ShortWriter(val valueVector: SmallIntVector) } private[arrow] class IntegerWriter(val valueVector: IntVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getInt(ordinal)) @@ -988,7 +1150,7 @@ private[arrow] class IntegerWriter(val valueVector: IntVector) } private[arrow] class LongWriter(val valueVector: BigIntVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getLong(ordinal)) @@ -1000,7 +1162,7 @@ private[arrow] class LongWriter(val valueVector: BigIntVector) } private[arrow] class FloatWriter(val valueVector: Float4Vector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getFloat(ordinal)) @@ -1012,7 +1174,7 @@ private[arrow] class FloatWriter(val valueVector: Float4Vector) } private[arrow] class DoubleWriter(val valueVector: Float8Vector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getDouble(ordinal)) @@ -1031,7 +1193,7 @@ private[arrow] class DoubleWriter(val valueVector: Float8Vector) * that `getDecimal` would build is only built for a value that does not fit, to fail as before. */ private[arrow] class DecimalWriter(val valueVector: DecimalVector, precision: Int, scale: Int) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { import ArrowFieldWriter.LittleEndian override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { @@ -1039,23 +1201,49 @@ private[arrow] class DecimalWriter(val valueVector: DecimalVector, precision: In setValueUnsafe(input, ordinal) } + // Spark's unsafe rows and arrays hold up to 18 digits as the unscaled long, and more as the + // unscaled bytes, which are read in place. An unsafe row's getDecimal does not check the long + // against the precision, but an unsafe array's does, so a long that does not fit is left to it. override protected def setValueUnsafe(input: SpecializedGetters, ordinal: Int): Unit = { - if (precision > Decimal.MAX_LONG_DIGITS && LittleEndian && - (input.isInstanceOf[UnsafeRow] || input.isInstanceOf[UnsafeArrayData])) { - val bytes = input.getBinary(ordinal) - if (putBigEndianIfFits(count, bytes, Platform.BYTE_ARRAY_OFFSET.toLong, bytes.length)) { - BitVectorHelper.setBit(valueVector.getValidityBuffer, count.toLong) - return - } - } - val decimal = input.getDecimal(ordinal, precision, scale) - if (decimal.changePrecision(precision, scale)) { - setUnscaled(count, decimal) + val written = LittleEndian && (input match { + case row: UnsafeRow if precision <= Decimal.MAX_LONG_DIGITS => + putLong(valueAddress, row.getLong(ordinal)) + true + case array: UnsafeArrayData if precision <= Decimal.MAX_LONG_DIGITS => + val unscaled = array.getLong(ordinal) + val fits = DecimalWriter.fitsPrecision(unscaled >> 63, unscaled, precision) + if (fits) { + putLong(valueAddress, unscaled) + } + fits + case row: UnsafeRow => + putUnsafeBytesIfFits(row.getBaseObject, row.getBaseOffset, row.getLong(ordinal)) + case array: UnsafeArrayData => + putUnsafeBytesIfFits(array.getBaseObject, array.getBaseOffset, array.getLong(ordinal)) + case _ => false + }) + if (written) { + BitVectorHelper.setBit(valueVector.getValidityBuffer, count.toLong) } else { - setNullUnsafe() + val decimal = input.getDecimal(ordinal, precision, scale) + if (decimal.changePrecision(precision, scale)) { + setUnscaled(count, decimal) + } else { + setNullUnsafe() + } } } + private def valueAddress: Long = + valueVector.getDataBufferAddress + count.toLong * DecimalVector.TYPE_WIDTH + + /** + * Writes the unscaled bytes that `offsetAndSize` locates in an unsafe row or array to value + * `count`, leaving validity alone, if they fit the precision. Returns whether it did. + */ + private def putUnsafeBytesIfFits(base: AnyRef, baseOffset: Long, offsetAndSize: Long): Boolean = + putBigEndianIfFits(count, base, baseOffset + (offsetAndSize >> 32), offsetAndSize.toInt) + /** Sets the value and validity of `index`, which must be within capacity. */ private def setUnscaled(index: Int, decimal: Decimal): Unit = { if (precision <= Decimal.MAX_LONG_DIGITS) { @@ -1082,7 +1270,7 @@ private[arrow] class DecimalWriter(val valueVector: DecimalVector, precision: In * whether it did. */ private def putBigEndianIfFits(index: Int, base: AnyRef, offset: Long, length: Int): Boolean = { - if (length == 0 || length > DecimalVector.TYPE_WIDTH) { + if (length <= 0 || length > DecimalVector.TYPE_WIDTH) { return false } var low = if (Platform.getByte(base, offset) < 0) -1L else 0L @@ -1208,37 +1396,26 @@ private[arrow] object DecimalWriter { } } -private[arrow] class StringWriter(val valueVector: VarCharVector) extends ArrowFieldWriter { +/** + * Writes strings or binaries. Spark's own vectors are copied in bulk, and its unsafe rows and + * arrays straight from their memory: both keep a value's offset and size in its fixed-length + * slot, so its bytes are copied without a `UTF8String` or Arrow's per-value setters. Other input + * goes through `setValue`. + */ +private[arrow] abstract class VariableWidthArrowFieldWriter extends ArrowFieldWriter { + import ArrowFieldWriter._ + + override def valueVector: BaseVariableWidthVector override def setNull(): Unit = { valueVector.setNull(count) } - override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { - val utf8 = input.getUTF8String(ordinal) - if (utf8.getBaseObject == null) { - val length = utf8.numBytes() - require(length >= 0, "String length must be non-negative") - valueVector.setValueLengthSafe(count, length) - - // Reservation can replace the buffer. Copy into its current address while Spark still owns - // the source bytes, without staging the off-heap payload in a JVM byte array. - val data = valueVector.getDataBuffer - val offset = valueVector.getStartOffset(count).toLong - require(offset >= 0 && offset + length <= data.capacity(), "Invalid Arrow string range") - utf8.writeToMemory(null, Math.addExact(data.memoryAddress(), offset)) - valueVector.setIndexDefined(count) - } else { - val utf8ByteBuffer = utf8.getByteBuffer - valueVector.setSafe(count, utf8ByteBuffer, utf8ByteBuffer.position(), utf8.numBytes()) - } - } - override def writeColumnSlice(input: ColumnVector, startRow: Int, numRows: Int): Unit = { input match { - case vector: WritableColumnVector if ArrowFieldWriter.isSparkVector(vector) => + case vector: WritableColumnVector if isSparkVector(vector) => if (vector.hasDictionary) { - ArrowFieldWriter.writeDictionaryVariableWidth( + writeDictionaryVariableWidth( valueVector, count, vector, @@ -1246,7 +1423,7 @@ private[arrow] class StringWriter(val valueVector: VarCharVector) extends ArrowF numRows, dictionaryCache) } else { - ArrowFieldWriter.writeVariableWidth(valueVector, count, vector, startRow, numRows) + writeVariableWidth(valueVector, count, vector, startRow, numRows) } count += numRows case _ => @@ -1262,6 +1439,102 @@ private[arrow] class StringWriter(val valueVector: VarCharVector) extends ArrowF super.reset() dictionaryCache.clear() } + + // The value is checked before Arrow grows anything, so one that cannot be written changes + // nothing. + override private[arrow] def writeUnsafeRowField(row: UnsafeRow, ordinal: Int): Unit = { + if (row.isNullAt(ordinal)) { + setNull() + } else { + val offsetAndSize = row.getLong(ordinal) + val start = valueVector.getStartOffset(valueVector.getLastSet + 1) + val length = (checkedEnd(start.toLong, offsetAndSize) - start).toInt + // Grows the buffers, and gives null rows written since the last value their empty offsets. + valueVector.setValueLengthSafe(count, length) + copyMemory( + row.getBaseObject, + row.getBaseOffset + (offsetAndSize >> 32), + valueVector.getDataBuffer.memoryAddress + start, + length.toLong) + BitVectorHelper.setBit(valueVector.getValidityBuffer, count.toLong) + } + count += 1 + } + + // Every element is checked before anything changes, as one value is, so the buffers grow once + // and the bytes are copied without further checks. + override private[arrow] def writeArrayElements(array: UnsafeArrayData): Unit = { + val numElements = array.numElements() + if (numElements == 0) { + return + } + val start = valueVector.getStartOffset(valueVector.getLastSet + 1) + var end = start.toLong + var i = 0 + while (i < numElements) { + if (!array.isNullAt(i)) { + end = checkedEnd(end, array.getLong(i)) + } + i += 1 + } + reserveValues(valueVector, count, numElements) + reserveData(valueVector, end) + val base = array.getBaseObject + val baseOffset = array.getBaseOffset + val data = valueVector.getDataBuffer.memoryAddress + val offsets = valueVector.getOffsetBuffer + var offset = start.toLong + i = 0 + while (i < numElements) { + if (!array.isNullAt(i)) { + val offsetAndSize = array.getLong(i) + val length = offsetAndSize.toInt + copyMemory(base, baseOffset + (offsetAndSize >> 32), data + offset, length.toLong) + offset += length + } + offsets.setInt((count + i + 1).toLong * BaseVariableWidthVector.OFFSET_WIDTH, offset.toInt) + i += 1 + } + writeArrayValidity(valueVector.getValidityBuffer, count, array) + valueVector.setLastSet(count + numElements - 1) + count += numElements + } + + /** The end of a value written at `start` with the size packed into `offsetAndSize`. */ + private def checkedEnd(start: Long, offsetAndSize: Long): Long = { + val length = offsetAndSize.toInt + if (length < 0) { + val kind = if (valueVector.isInstanceOf[VarCharVector]) "String" else "Binary" + throw new IllegalArgumentException(s"$kind length must be non-negative") + } + val end = start + length + checkDataEnd(end) + end + } +} + +private[arrow] class StringWriter(val valueVector: VarCharVector) + extends VariableWidthArrowFieldWriter { + + override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { + val utf8 = input.getUTF8String(ordinal) + if (utf8.getBaseObject == null) { + val length = utf8.numBytes() + require(length >= 0, "String length must be non-negative") + valueVector.setValueLengthSafe(count, length) + + // Reservation can replace the buffer. Copy into its current address while Spark still owns + // the source bytes, without staging the off-heap payload in a JVM byte array. + val data = valueVector.getDataBuffer + val offset = valueVector.getStartOffset(count).toLong + require(offset >= 0 && offset + length <= data.capacity(), "Invalid Arrow string range") + utf8.writeToMemory(null, Math.addExact(data.memoryAddress(), offset)) + valueVector.setIndexDefined(count) + } else { + val utf8ByteBuffer = utf8.getByteBuffer + valueVector.setSafe(count, utf8ByteBuffer, utf8ByteBuffer.position(), utf8.numBytes()) + } + } } private[arrow] class LargeStringWriter(val valueVector: LargeVarCharVector) @@ -1279,45 +1552,13 @@ private[arrow] class LargeStringWriter(val valueVector: LargeVarCharVector) } } -private[arrow] class BinaryWriter(val valueVector: VarBinaryVector) extends ArrowFieldWriter { - - override def setNull(): Unit = { - valueVector.setNull(count) - } +private[arrow] class BinaryWriter(val valueVector: VarBinaryVector) + extends VariableWidthArrowFieldWriter { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { val bytes = input.getBinary(ordinal) valueVector.setSafe(count, bytes, 0, bytes.length) } - - override def writeColumnSlice(input: ColumnVector, startRow: Int, numRows: Int): Unit = { - input match { - case vector: WritableColumnVector if ArrowFieldWriter.isSparkVector(vector) => - if (vector.hasDictionary) { - ArrowFieldWriter.writeDictionaryVariableWidth( - valueVector, - count, - vector, - startRow, - numRows, - dictionaryCache) - } else { - ArrowFieldWriter.writeVariableWidth(valueVector, count, vector, startRow, numRows) - } - count += numRows - case _ => - super.writeColumnSlice(input, startRow, numRows) - } - } - - private val dictionaryCache = new DictionaryCache - - override private[arrow] def startInputBatch(): Unit = dictionaryCache.clear() - - override def reset(): Unit = { - super.reset() - dictionaryCache.clear() - } } private[arrow] class LargeBinaryWriter(val valueVector: LargeVarBinaryVector) @@ -1334,7 +1575,7 @@ private[arrow] class LargeBinaryWriter(val valueVector: LargeVarBinaryVector) } private[arrow] class DateWriter(val valueVector: DateDayVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getInt(ordinal)) @@ -1346,7 +1587,7 @@ private[arrow] class DateWriter(val valueVector: DateDayVector) } private[arrow] class TimestampWriter(val valueVector: TimeStampMicroTZVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getLong(ordinal)) @@ -1358,7 +1599,7 @@ private[arrow] class TimestampWriter(val valueVector: TimeStampMicroTZVector) } private[arrow] class TimestampNTZWriter(val valueVector: TimeStampMicroVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getLong(ordinal)) @@ -1370,7 +1611,7 @@ private[arrow] class TimestampNTZWriter(val valueVector: TimeStampMicroVector) } private[arrow] class TimeNanoWriter(val valueVector: TimeNanoVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getLong(ordinal)) @@ -1387,13 +1628,49 @@ private[arrow] class ArrayWriter(val valueVector: ListVector, val elementWriter: override def setNull(): Unit = {} override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { - val array = input.getArray(ordinal) + input.getArray(ordinal) match { + case unsafe: UnsafeArrayData => + writeElements(unsafe) + case array => + val numElements = array.numElements() + valueVector.startNewValue(count) + var i = 0 + while (i < numElements) { + elementWriter.write(array, i) + i += 1 + } + valueVector.endValue(count, numElements) + } + } + + // Each unsafe array gets a view of its own from Spark's getter. A view reused across values + // would keep the last row it read reachable after that row's values were copied. + override private[arrow] def writeUnsafeRowField(row: UnsafeRow, ordinal: Int): Unit = { + if (row.isNullAt(ordinal)) { + setNull() + } else { + writeElements(row.getArray(ordinal)) + } + count += 1 + } + + override private[arrow] def writeArrayElements(array: UnsafeArrayData): Unit = { + val numElements = array.numElements() var i = 0 - valueVector.startNewValue(count) - while (i < array.numElements()) { - elementWriter.write(array, i) + while (i < numElements) { + if (array.isNullAt(i)) { + setNull() + } else { + writeElements(array.getArray(i)) + } + count += 1 i += 1 } + } + + private def writeElements(array: UnsafeArrayData): Unit = { + valueVector.startNewValue(count) + elementWriter.writeArrayElements(array) valueVector.endValue(count, array.numElements()) } @@ -1449,11 +1726,48 @@ private[arrow] class StructWriter( } override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { - val struct = input.getStruct(ordinal, children.length) + input.getStruct(ordinal, children.length) match { + case unsafe: UnsafeRow => + writeFields(unsafe) + case struct => + var i = 0 + valueVector.setIndexDefined(count) + while (i < struct.numFields) { + children(i).write(struct, i) + i += 1 + } + } + } + + // Each unsafe struct gets a view of its own, as in ArrayWriter. + override private[arrow] def writeUnsafeRowField(row: UnsafeRow, ordinal: Int): Unit = { + if (row.isNullAt(ordinal)) { + setNull() + } else { + writeFields(row.getStruct(ordinal, children.length)) + } + count += 1 + } + + override private[arrow] def writeArrayElements(array: UnsafeArrayData): Unit = { + val numElements = array.numElements() var i = 0 + while (i < numElements) { + if (array.isNullAt(i)) { + setNull() + } else { + writeFields(array.getStruct(i, children.length)) + } + count += 1 + i += 1 + } + } + + private def writeFields(struct: UnsafeRow): Unit = { valueVector.setIndexDefined(count) - while (i < struct.numFields) { - children(i).write(struct, i) + var i = 0 + while (i < children.length) { + children(i).writeUnsafeRowField(struct, i) i += 1 } } @@ -1524,19 +1838,65 @@ private[arrow] class MapWriter( override def setNull(): Unit = {} override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { - val map = input.getMap(ordinal) - valueVector.startNewValue(count) - val keys = map.keyArray() - val values = map.valueArray() + input.getMap(ordinal) match { + case unsafe: UnsafeMapData => + writeEntries(unsafe) + case map => + val numElements = map.numElements() + valueVector.startNewValue(count) + val keys = map.keyArray() + val values = map.valueArray() + var i = 0 + while (i < numElements) { + structVector.setIndexDefined(keyWriter.count) + keyWriter.write(keys, i) + valueWriter.write(values, i) + i += 1 + } + valueVector.endValue(count, numElements) + } + } + + // Each unsafe map gets a view of its own, as in ArrayWriter. + override private[arrow] def writeUnsafeRowField(row: UnsafeRow, ordinal: Int): Unit = { + if (row.isNullAt(ordinal)) { + setNull() + } else { + writeEntries(row.getMap(ordinal)) + } + count += 1 + } + + override private[arrow] def writeArrayElements(array: UnsafeArrayData): Unit = { + val numElements = array.numElements() var i = 0 - while (i < map.numElements()) { - structVector.setIndexDefined(keyWriter.count) - keyWriter.write(keys, i) - valueWriter.write(values, i) + while (i < numElements) { + if (array.isNullAt(i)) { + setNull() + } else { + writeEntries(array.getMap(i)) + } + count += 1 i += 1 } + } + + private def writeEntries(map: UnsafeMapData): Unit = { + val numElements = map.numElements() + valueVector.startNewValue(count) + if (numElements > 0) { + setEntriesValid(keyWriter.count, numElements) + keyWriter.writeArrayElements(map.keyArray()) + valueWriter.writeArrayElements(map.valueArray()) + } + valueVector.endValue(count, numElements) + } - valueVector.endValue(count, map.numElements()) + /** Marks entries `[start, start + numEntries)` valid. */ + private def setEntriesValid(start: Int, numEntries: Int): Unit = { + // Grows the validity buffer to the last entry. + structVector.setIndexDefined(start + numEntries - 1) + ArrowFieldWriter.setValid(structVector.getValidityBuffer, start, numEntries) } // Like ArrayWriter.writeColumnSlice, with the keys and values in Spark's two child vectors. @@ -1547,9 +1907,7 @@ private[arrow] class MapWriter( val numEntries = ArrowFieldWriter.writeListOffsets(valueVector, count, vector, startRow, numRows) if (numEntries > 0) { - // Grows the entries' validity buffer, then marks every entry valid. - structVector.setIndexDefined(entryStart + numEntries - 1) - ArrowFieldWriter.setValid(structVector.getValidityBuffer, entryStart, numEntries) + setEntriesValid(entryStart, numEntries) } val keys = vector.getChild(0) val values = vector.getChild(1) @@ -1600,7 +1958,7 @@ private[arrow] class NullWriter(val valueVector: NullVector) extends ArrowFieldW } private[arrow] class IntervalYearWriter(val valueVector: IntervalYearVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getInt(ordinal)) @@ -1612,7 +1970,7 @@ private[arrow] class IntervalYearWriter(val valueVector: IntervalYearVector) } private[arrow] class DurationWriter(val valueVector: DurationVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { valueVector.setSafe(count, input.getLong(ordinal)) @@ -1624,7 +1982,7 @@ private[arrow] class DurationWriter(val valueVector: DurationVector) } private[arrow] class IntervalMonthDayNanoWriter(val valueVector: IntervalMonthDayNanoVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { val ci = input.getInterval(ordinal) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala index 19532cefa14..316afe68a22 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometNativeArrowSource.scala @@ -19,11 +19,15 @@ package org.apache.spark.sql.comet.execution.arrow +import java.util.concurrent.ConcurrentHashMap + import scala.jdk.CollectionConverters._ import scala.reflect.ClassTag import org.apache.arrow.c.{ArrowArrayStream, Data} import org.apache.arrow.memory.BufferAllocator +import org.apache.arrow.vector.VectorSchemaRoot +import org.apache.arrow.vector.dictionary.Dictionary import org.apache.arrow.vector.ipc.ArrowReader import org.apache.arrow.vector.types.pojo.{Field, FieldType, Schema} import org.apache.spark.TaskContext @@ -233,16 +237,19 @@ object CometArrowStream extends Logging { * the JNI global ref that arrow-java holds on the stream's private data. Taking the stream * leaves a null release callback in the JVM's struct, and arrow-java's `release` skips a null * callback, so a stream native took is still released only once. + * + * The reader is wrapped so that what it throws while native pulls a batch stays available to + * [[inputFailure]] until the task completes. */ def stream( name: String, readerFactory: BufferAllocator => ArrowReader): Iterator[ArrowArrayStream] = { val context = TaskContext.get() val allocator = CometArrowAllocator.newChildAllocator(name, 0, Long.MaxValue) - var reader: ArrowReader = null + var reader: FailureRecordingReader = null var arrowStream: ArrowArrayStream = null try { - reader = readerFactory(allocator) + reader = new FailureRecordingReader(allocator, readerFactory(allocator)) arrowStream = ArrowArrayStream.allocateNew(allocator) Data.exportArrayStream(allocator, reader, arrowStream) } catch { @@ -256,7 +263,9 @@ object CometArrowStream extends Logging { } if (context != null) { val streamRef = arrowStream + exportedReaders.put(streamRef, reader) context.addTaskCompletionListener[Unit] { _ => + exportedReaders.remove(streamRef) // Release before close: `close` frees the struct that `release` reads. streamRef.release() streamRef.close() @@ -266,6 +275,59 @@ object CometArrowStream extends Logging { Iterator.single(arrowStream) } + /** + * The reader behind each stream exported by [[stream]] whose task has not completed yet, keyed + * by the stream object that its consumer holds as an input. + */ + private val exportedReaders = new ConcurrentHashMap[ArrowArrayStream, FailureRecordingReader]() + + /** + * What the reader behind `stream` threw while native pulled a batch from it, if anything. Arrow + * Java's exported stream catches that throwable and hands native only its text, so the native + * plan fails with a `CometNativeException` built from the text. The throwable itself is what + * the task would have thrown without Comet, so the plan's consumer rethrows it instead. + */ + def inputFailure(stream: ArrowArrayStream): Option[Throwable] = + Option(exportedReaders.get(stream)).flatMap(reader => Option(reader.failure)) + + /** + * Forwards everything to `reader`, keeping the first throwable its `loadNextBatch` throws for + * [[inputFailure]]. Its own `ArrowReader` state is never initialized: every method that the + * stream exporter calls is forwarded, and so is `close`. + */ + private final class FailureRecordingReader(allocator: BufferAllocator, reader: ArrowReader) + extends ArrowReader(allocator) { + + @volatile var failure: Throwable = _ + + override def loadNextBatch(): Boolean = + try reader.loadNextBatch() + catch { + case t: Throwable => + if (failure == null) failure = t + throw t + } + + override def getVectorSchemaRoot(): VectorSchemaRoot = reader.getVectorSchemaRoot() + + override def getDictionaryVectors(): java.util.Map[java.lang.Long, Dictionary] = + reader.getDictionaryVectors() + + override def lookup(id: Long): Dictionary = reader.lookup(id) + + override def getDictionaryIds(): java.util.Set[java.lang.Long] = reader.getDictionaryIds() + + override def bytesRead(): Long = reader.bytesRead() + + override def close(): Unit = reader.close() + + override def close(closeReadSource: Boolean): Unit = reader.close(closeReadSource) + + override protected def closeReadSource(): Unit = () + + override protected def readSchema(): Schema = reader.getVectorSchemaRoot().getSchema + } + /** Close a resource, swallowing any error, for use in rollback paths. Null-safe. */ private def closeQuietly(c: AutoCloseable): Unit = { if (c != null) { diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index 3216b6a4301..c5e5c7641e5 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -21,6 +21,7 @@ package org.apache.spark.sql.comet.execution.shuffle import java.util.function.Supplier +import scala.collection.mutable.ListBuffer import scala.concurrent.Future import scala.jdk.CollectionConverters._ @@ -214,17 +215,18 @@ case class CometShuffleExchangeExec( val nativeChild = child.asInstanceOf[CometNativeExec] // RangePartitioner needs real rows for sampling. Reuse the precomputed context so we // don't re-walk the SparkPlan tree or re-broadcast the encryption Hadoop conf. - val samplingRDD: Option[RDD[ColumnarBatch]] = outputPartitioning match { + val samplingRows: Option[RDD[InternalRow]] = outputPartitioning match { case _: RangePartitioning => Some( - nativeChild.executeColumnarWithContext( - ctx, - nativeChildMetricNode.withoutAggregateMetrics(nativeChild))) + CometShuffleExchangeExec.rowsOf( + nativeChild.executeColumnarWithContext( + ctx, + nativeChildMetricNode.withoutAggregateMetrics(nativeChild)))) case _ => None } CometShuffleExchangeExec.prepareNativeShuffleDependency( inputRDD.asInstanceOf[CometNativeShuffleInputRDD], - samplingRDD, + samplingRows, child.output, outputPartitioning, serializer, @@ -236,12 +238,19 @@ case class CometShuffleExchangeExec( positionalRoundRobin)) case None => child match { - case _: CometNativeArrowSource => + case source: CometNativeArrowSource => CometShuffleExchangeExec.prepareArrowStreamShuffleDependency( inputRDD.asInstanceOf[RDD[ArrowArrayStream]], - // The range partitioner samples rows, so it needs them as batches. outputPartitioning match { - case _: RangePartitioning => Some(child.executeColumnar()) + // The range partitioner samples only the sort keys, so a conversion of rows + // gives it the rows it reads, rather than all of their columns in Arrow. + case _: RangePartitioning => + Some(source match { + case conversion: CometSparkToColumnarExec + if !conversion.child.supportsColumnar => + conversion.child.execute() + case _ => CometShuffleExchangeExec.rowsOf(child.executeColumnar()) + }) case _ => None }, child.output, @@ -547,6 +556,48 @@ object CometShuffleExchangeExec None } + /** + * Whether a shuffle that would use the JVM columnar shuffle, because its child is a Spark + * row-based plan, can convert the child's rows to Arrow with `CometSparkToColumnarExec` and use + * native shuffle instead. See [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. The + * native checks are those that [[shuffleSupported]] would make for the converted plan, so + * `CometExecRule` only converts a shuffle that will become native. Like [[shuffleSupported]], + * tags the node when no Comet shuffle can take it. + */ + def convertsInputForNativeShuffle(s: ShuffleExchangeExec): Boolean = + CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.get(s.conf) && + !isCometPlan(s.child) && + // This also rules out Celeborn, which has no JVM columnar shuffle, and calendar intervals. + // Arrow holds the time part of an interval in nanoseconds, so the conversion would + // overflow on one with more microseconds than that can hold. + shuffleSupported(s).contains(CometColumnarShuffle) && + CometSparkToColumnarExec.isSchemaSupported(s.child.schema, ListBuffer.empty) && + nativeShuffleFailureReasons(s).isEmpty && + !hashesDifferentlyFromSpark(s) + + /** + * Whether the shuffle hashes into more than one partition a key that native shuffle, reading + * converted rows, may put in a different partition from Spark's partitioner. Native shuffle + * must not take over such a shuffle from the JVM columnar shuffle: a join with an input that is + * still partitioned by Spark would put matching keys in different partitions. The keys are: + * + * - A decimal wider than 18 digits, which native shuffle hashes differently (#5994). + * - A string, or any value computed from one, such as `hash(s)`. Spark reads the string's + * bytes as they are, but the import of the converted batch into native replaces invalid + * UTF-8 before native shuffle evaluates the key. + * + * TODO: allow wide decimals once native hashing matches Spark for them. + */ + private def hashesDifferentlyFromSpark(s: ShuffleExchangeExec): Boolean = + s.outputPartitioning match { + case HashPartitioning(expressions, numPartitions) => + numPartitions > 1 && expressions.exists { key => + key.dataType.existsRecursively(DecimalType.isByteArrayDecimalType) || + key.exists(_.dataType.existsRecursively(_.isInstanceOf[StringType])) + } + case _ => false + } + /** * Whether the stage feeding a shuffle starts at a typed Dataset conversion (see * [[CometConf.COMET_CONVERT_FROM_TYPED_DATASET_ENABLED]]). `CometExecRule` decides the shuffle @@ -862,6 +913,10 @@ object CometShuffleExchangeExec } } + /** The rows of `batches`, for the range partitioner to sample. */ + private def rowsOf(batches: RDD[ColumnarBatch]): RDD[InternalRow] = + batches.mapPartitionsInternal(_.flatMap(_.rowIterator().asScala)) + /** * Build a Comet native shuffle dependency around an existing `RDD[ColumnarBatch]` of real * batches. Used by [[org.apache.spark.sql.comet.CometCollectLimitExec]] and @@ -891,7 +946,7 @@ object CometShuffleExchangeExec "ShuffleWriterInput") prepareArrowStreamShuffleDependency( streamRDD, - Some(rdd), + Some(rowsOf(rdd)), outputAttributes, outputPartitioning, serializer, @@ -900,12 +955,12 @@ object CometShuffleExchangeExec /** * [[prepareShuffleDependency]] for input that native reads as one Arrow stream per partition, - * such as a [[CometNativeArrowSource]]'s. `samplingRDD` gives the same rows as batches, and is - * only required for [[RangePartitioning]]. + * such as a [[CometNativeArrowSource]]'s. `samplingRows` gives the same rows, and is only + * required for [[RangePartitioning]]. */ private def prepareArrowStreamShuffleDependency( streamRDD: RDD[ArrowArrayStream], - samplingRDD: Option[RDD[ColumnarBatch]], + samplingRows: Option[RDD[InternalRow]], outputAttributes: Seq[Attribute], outputPartitioning: Partitioning, serializer: Serializer, @@ -946,7 +1001,7 @@ object CometShuffleExchangeExec // is `shuffleWriterMetrics` at the root with one empty leaf for the Scan child. prepareNativeShuffleDependency( thinRDD, - samplingRDD, + samplingRows, outputAttributes, outputPartitioning, serializer, @@ -965,13 +1020,14 @@ object CometShuffleExchangeExec * @param thinRDD * scheduling-anchor RDD whose `compute` returns a [[CometNativeShuffleInputIterator]]; * produces no batches itself. - * @param samplingRDD - * regular columnar execution of the child, only required for [[RangePartitioning]] (sampling - * needs real rows). `None` for hash / single / round-robin. + * @param samplingRows + * the child's rows, such as those of a regular columnar execution of the child, only required + * for [[RangePartitioning]] (sampling needs real rows). `None` for hash / single / + * round-robin. */ def prepareNativeShuffleDependency( thinRDD: CometNativeShuffleInputRDD, - samplingRDD: Option[RDD[ColumnarBatch]], + samplingRows: Option[RDD[InternalRow]], outputAttributes: Seq[Attribute], outputPartitioning: Partitioning, serializer: Serializer, @@ -1011,11 +1067,11 @@ object CometShuffleExchangeExec // ShuffleExchangeExec::prepareShuffleDependency val (partitioner, rangePartitionBounds) = outputPartitioning match { case rangePartitioning: RangePartitioning => - // Sampling needs real rows; use the dedicated samplingRDD (a regular columnar execution - // of the child). The thin RDD itself yields nothing. - val samplingInput = samplingRDD.getOrElse( + // Sampling needs real rows; use the dedicated samplingRows (the child's rows). The thin + // RDD itself yields nothing. + val samplingInput = samplingRows.getOrElse( throw new IllegalStateException( - "RangePartitioning requires a samplingRDD on the native-shuffle path")) + "RangePartitioning requires samplingRows on the native-shuffle path")) // Extract only fields used for sorting to avoid collecting large fields that does not // affect sorting result when deciding partition bounds in RangePartitioner val rddForSampling = samplingInput.mapPartitionsInternal { iter => @@ -1025,12 +1081,7 @@ object CometShuffleExchangeExec // Internally, RangePartitioner runs a job on the RDD that samples keys to compute // partition bounds. To get accurate samples, we need to copy the mutable keys. - iter.flatMap { batch => - val rowIter = batch.rowIterator().asScala - rowIter.map { row => - mutablePair.update(projection(row).copy(), null) - } - } + iter.map(row => mutablePair.update(projection(row).copy(), null)) } // Construct ordering on extracted sort key. diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index 909539bf2b6..1516fdd9209 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -661,6 +661,30 @@ abstract class CometExec extends CometPlan { /** The original Spark operator from which this Comet operator is converted from */ def originalPlan: SparkPlan + /** + * Rebuilds the Spark operator represented by this Comet operator with reverted children. + * + * Operators whose Comet representation changes the Spark plan shape or whose live state differs + * from the original Spark plan must override this method. + */ + def sparkFallback(newChildren: Seq[SparkPlan]): SparkPlan = { + val sparkPlan = originalPlan + if (sparkPlan == null) { + throw new CometExec.InvalidSparkFallbackException( + s"${getClass.getSimpleName} has no original Spark plan") + } + if (newChildren.exists(_ eq sparkPlan)) { + throw new CometExec.InvalidSparkFallbackException( + s"${getClass.getSimpleName} aliases its original Spark plan with a child") + } + if (sparkPlan.children.size != newChildren.size) { + throw new CometExec.InvalidSparkFallbackException( + s"${getClass.getSimpleName} cannot restore ${sparkPlan.getClass.getSimpleName}: " + + s"expected ${sparkPlan.children.size} children but received ${newChildren.size}") + } + sparkPlan.withNewChildren(newChildren) + } + /** Comet always support columnar execution */ override def supportsColumnar: Boolean = true @@ -711,6 +735,9 @@ abstract class CometExec extends CometPlan { } object CometExec { + final class InvalidSparkFallbackException(message: String) + extends IllegalArgumentException(message) + // An unique id for each CometExecIterator, used to identify the native query execution. private val curId = new java.util.concurrent.atomic.AtomicLong() @@ -1714,6 +1741,13 @@ object CometExpandExec extends CometOperatorSerde[ExpandExec] { op: ExpandExec, builder: Operator.Builder, childOp: OperatorOuterClass.Operator*): Option[OperatorOuterClass.Operator] = { + // The native Expand regroups the flattened expressions by their count per projection, so a + // projection with no columns, which column pruning leaves when nothing above the Expand reads + // its output, would reach it as no projections at all. + if (op.projections.isEmpty || op.projections.head.isEmpty) { + withFallbackReason(op, "Expand without output columns is not supported") + return None + } val projExprs = op.projections.flatMap(_.map(e => exprToProto(e, op.child.output))) if (projExprs.forall(_.isDefined) && childOp.nonEmpty) { @@ -1866,12 +1900,8 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { return Unsupported(Some(s"Unsupported generator: ${op.generator.nodeName}")) } op.generator.children.head.dataType match { - case _: ArrayType => + case _: ArrayType | _: MapType => Compatible() - case _: MapType => - // TODO add support for map types - // https://github.com/apache/datafusion-comet/issues/2837 - Unsupported(Some("Comet only supports explode/explode_outer for arrays, not maps")) case other => Unsupported(Some(s"Unsupported data type: $other")) } diff --git a/spark/src/main/spark-3.4/org/apache/comet/iceberg/IcebergDeltaLogicalShim.scala b/spark/src/main/spark-3.4/org/apache/comet/iceberg/IcebergDeltaLogicalShim.scala new file mode 100644 index 00000000000..afbf7d60b70 --- /dev/null +++ b/spark/src/main/spark-3.4/org/apache/comet/iceberg/IcebergDeltaLogicalShim.scala @@ -0,0 +1,27 @@ +/* + * 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.iceberg + +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan + +/** WriteDelta interception starts with Spark 3.5. */ +private[iceberg] object IcebergDeltaLogicalShim extends IcebergDeltaLogicalShimApi { + override def extract(plan: LogicalPlan): Option[DeltaLogicalFields] = None +} diff --git a/spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala b/spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala index 4fd3606203b..f47c7682389 100644 --- a/spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala +++ b/spark/src/main/spark-3.4/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala @@ -294,6 +294,12 @@ trait ShimSparkErrorConverter { messageParameters = params.map { case (k, v) => (k, v.toString) }, cause = null)) + case "InvalidUrl" => + Some( + QueryExecutionErrors.invalidUrlError( + UTF8String.fromString(params("url").toString), + new java.net.URISyntaxException(params("url").toString, "Invalid URL"))) + case "DatatypeCannotOrder" => // orderedOperationUnsupportedByDataTypeError takes DataType in Spark 3.4, not String Some( diff --git a/spark/src/main/spark-3.5/org/apache/comet/iceberg/IcebergDeltaLogicalShim.scala b/spark/src/main/spark-3.5/org/apache/comet/iceberg/IcebergDeltaLogicalShim.scala new file mode 100644 index 00000000000..d85d0c19fd8 --- /dev/null +++ b/spark/src/main/spark-3.5/org/apache/comet/iceberg/IcebergDeltaLogicalShim.scala @@ -0,0 +1,47 @@ +/* + * 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.iceberg + +import org.apache.spark.sql.catalyst.analysis.NamedRelation +import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, WriteDelta} +import org.apache.spark.sql.connector.write.RowLevelOperation.Command.{DELETE, MERGE, UPDATE} + +/** Spark 3.5 exposes the WriteDelta fields directly; keep extraction fully typed. */ +private[iceberg] object IcebergDeltaLogicalShim extends IcebergDeltaLogicalShimApi { + override def extract(plan: LogicalPlan): Option[DeltaLogicalFields] = plan match { + case delta: WriteDelta => + val command = delta.operation.command match { + case DELETE => Some(DeltaDelete) + case UPDATE => Some(DeltaUpdate) + case MERGE => Some(DeltaMerge) + case _ => None + } + command.map { value => + val relation: NamedRelation = delta.originalTable + DeltaLogicalFields( + query = delta.query, + originalTable = relation, + projections = delta.projections, + write = delta.write, + command = Some(value)) + } + case _ => None + } +} diff --git a/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala b/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala index 39a76f386c5..ee4dbea9725 100644 --- a/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala +++ b/spark/src/main/spark-3.5/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala @@ -291,6 +291,12 @@ trait ShimSparkErrorConverter { params("groupCount").toString.toInt, params("groupIndex").toString.toInt)) + case "InvalidUrl" => + Some( + QueryExecutionErrors.invalidUrlError( + UTF8String.fromString(params("url").toString), + new java.net.URISyntaxException(params("url").toString, "Invalid URL"))) + case "DatatypeCannotOrder" => Some( QueryExecutionErrors.orderedOperationUnsupportedByDataTypeError( diff --git a/spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergDeltaWriterShim.scala b/spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergDeltaWriterShim.scala new file mode 100644 index 00000000000..c652eb3558a --- /dev/null +++ b/spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergDeltaWriterShim.scala @@ -0,0 +1,65 @@ +/* + * 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.iceberg + +import org.apache.spark.SparkException +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.util.RowDeltaUtils.{DELETE_OPERATION, INSERT_OPERATION, UPDATE_OPERATION} +import org.apache.spark.sql.connector.write.DeltaWriter + +/** Spark 3.x supports DELETE, UPDATE, and INSERT operation codes in WriteDelta. */ +object IcebergDeltaWriterShim { + val OperationCodes: DeltaOperationCodes = + DeltaOperationCodes( + delete = DELETE_OPERATION, + update = UPDATE_OPERATION, + insert = INSERT_OPERATION, + reinsert = None) + + def writeOperation( + writer: DeltaWriter[InternalRow], + operation: Int, + row: InternalRow, + info: WriteDeltaDispatchInfo): Unit = { + def metadata(): InternalRow = info.metadataProjection.map { projection => + projection.project(row) + projection + }.orNull + + operation match { + case code if code == info.operationCodes.delete => + info.rowIdProjection.project(row) + writer.delete(metadata(), info.rowIdProjection) + case code if code == info.operationCodes.update => + val data = info.rowProjection.getOrElse( + throw new IllegalArgumentException("UPDATE operation has no row projection")) + data.project(row) + info.rowIdProjection.project(row) + writer.update(metadata(), info.rowIdProjection, data) + case code if code == info.operationCodes.insert => + val data = info.rowProjection.getOrElse( + throw new IllegalArgumentException("INSERT operation has no row projection")) + data.project(row) + writer.insert(data) + case _ => + throw new SparkException(s"Unexpected operation ID: $operation") + } + } +} diff --git a/spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala b/spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala new file mode 100644 index 00000000000..d4d145fc4ad --- /dev/null +++ b/spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala @@ -0,0 +1,30 @@ +/* + * 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.iceberg + +import org.apache.spark.SparkContext +import org.apache.spark.sql.execution.metric.SQLMetric + +/** Retry-aware row-level metrics are added in Spark 4.2. */ +object IcebergSemanticMetricsShim { + def deltaMetrics(sc: SparkContext, command: Option[DeltaCommand]): Map[String, SQLMetric] = + Map.empty + def value(metric: SQLMetric): Long = metric.value +} diff --git a/spark/src/main/spark-3.x/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala b/spark/src/main/spark-3.x/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala index d6a78ca2ea4..135384f5d9f 100644 --- a/spark/src/main/spark-3.x/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala +++ b/spark/src/main/spark-3.x/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala @@ -22,12 +22,15 @@ package org.apache.spark.sql.comet import org.apache.spark.sql.connector.write.{BatchWrite, WriterCommitMessage} import org.apache.spark.sql.execution.SparkPlan +import org.apache.comet.iceberg.DeltaCommand + /** The `BatchWrite.commit(messages, summary)` overload only exists on Spark 4.1+. */ private[comet] object IcebergWriteSummaryShim { def commit( batchWrite: BatchWrite, messages: Array[WriterCommitMessage], - query: SparkPlan): Unit = { + query: SparkPlan, + command: Option[DeltaCommand] = None): Unit = { batchWrite.commit(messages) } } diff --git a/spark/src/main/spark-4.0/org/apache/comet/iceberg/IcebergDeltaLogicalFieldsShim.scala b/spark/src/main/spark-4.0/org/apache/comet/iceberg/IcebergDeltaLogicalFieldsShim.scala new file mode 100644 index 00000000000..ea2ceba1823 --- /dev/null +++ b/spark/src/main/spark-4.0/org/apache/comet/iceberg/IcebergDeltaLogicalFieldsShim.scala @@ -0,0 +1,47 @@ +/* + * 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.iceberg + +import org.apache.spark.sql.catalyst.analysis.NamedRelation +import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, WriteDelta} +import org.apache.spark.sql.connector.write.RowLevelOperation.Command.{DELETE, MERGE, UPDATE} + +/** Spark 4.0 exposes the WriteDelta fields directly; keep extraction fully typed. */ +private[iceberg] object IcebergDeltaLogicalFieldsShim extends IcebergDeltaLogicalFieldsShimApi { + override def extract(plan: LogicalPlan): Option[DeltaLogicalFields] = plan match { + case delta: WriteDelta => + val command = delta.operation.command match { + case DELETE => Some(DeltaDelete) + case UPDATE => Some(DeltaUpdate) + case MERGE => Some(DeltaMerge) + case _ => None + } + command.map { value => + val relation: NamedRelation = delta.originalTable + DeltaLogicalFields( + query = delta.query, + originalTable = relation, + projections = delta.projections, + write = delta.write, + command = Some(value)) + } + case _ => None + } +} diff --git a/spark/src/main/spark-4.0/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala b/spark/src/main/spark-4.0/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala new file mode 100644 index 00000000000..d4d145fc4ad --- /dev/null +++ b/spark/src/main/spark-4.0/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala @@ -0,0 +1,30 @@ +/* + * 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.iceberg + +import org.apache.spark.SparkContext +import org.apache.spark.sql.execution.metric.SQLMetric + +/** Retry-aware row-level metrics are added in Spark 4.2. */ +object IcebergSemanticMetricsShim { + def deltaMetrics(sc: SparkContext, command: Option[DeltaCommand]): Map[String, SQLMetric] = + Map.empty + def value(metric: SQLMetric): Long = metric.value +} diff --git a/spark/src/main/spark-4.0/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala b/spark/src/main/spark-4.0/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala index d6a78ca2ea4..135384f5d9f 100644 --- a/spark/src/main/spark-4.0/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala +++ b/spark/src/main/spark-4.0/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala @@ -22,12 +22,15 @@ package org.apache.spark.sql.comet import org.apache.spark.sql.connector.write.{BatchWrite, WriterCommitMessage} import org.apache.spark.sql.execution.SparkPlan +import org.apache.comet.iceberg.DeltaCommand + /** The `BatchWrite.commit(messages, summary)` overload only exists on Spark 4.1+. */ private[comet] object IcebergWriteSummaryShim { def commit( batchWrite: BatchWrite, messages: Array[WriterCommitMessage], - query: SparkPlan): Unit = { + query: SparkPlan, + command: Option[DeltaCommand] = None): Unit = { batchWrite.commit(messages) } } diff --git a/spark/src/main/spark-4.1+/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala b/spark/src/main/spark-4.1+/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala index f3ea21e3171..af6c834bde1 100644 --- a/spark/src/main/spark-4.1+/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala +++ b/spark/src/main/spark-4.1+/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala @@ -24,7 +24,7 @@ import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.datasources.v2.MergeRowsExec -import org.apache.comet.shims.MergeRowsMetricsShim +import org.apache.comet.iceberg.{DeltaCommand, IcebergSemanticMetricsShim} /** * Spark 4.1 collects MERGE metrics from the executed plan and hands them to @@ -34,27 +34,30 @@ private[comet] object IcebergWriteSummaryShim extends AdaptiveSparkPlanHelper { def commit( batchWrite: BatchWrite, messages: Array[WriterCommitMessage], - query: SparkPlan): Unit = { - collectFirst(query) { - case m: MergeRowsExec => m.metrics - case m: CometMergeRowsExec => m.metrics - } match { - case Some(metrics) => - def metricValue(name: String): Long = - metrics.get(name).map(MergeRowsMetricsShim.value).getOrElse(-1L) - batchWrite.commit( - messages, - MergeSummaryImpl( - metricValue("numTargetRowsCopied"), - metricValue("numTargetRowsDeleted"), - metricValue("numTargetRowsUpdated"), - metricValue("numTargetRowsInserted"), - metricValue("numTargetRowsMatchedUpdated"), - metricValue("numTargetRowsMatchedDeleted"), - metricValue("numTargetRowsNotMatchedBySourceUpdated"), - metricValue("numTargetRowsNotMatchedBySourceDeleted"))) - case None => - batchWrite.commit(messages) + query: SparkPlan, + command: Option[DeltaCommand] = None): Unit = { + if (!IcebergDeltaWriteSummaryShim.commit(batchWrite, messages, query, command)) { + collectFirst(query) { + case m: MergeRowsExec => m.metrics + case m: CometMergeRowsExec => m.metrics + } match { + case Some(metrics) => + def metricValue(name: String): Long = + metrics.get(name).map(IcebergSemanticMetricsShim.value).getOrElse(-1L) + batchWrite.commit( + messages, + MergeSummaryImpl( + metricValue("numTargetRowsCopied"), + metricValue("numTargetRowsDeleted"), + metricValue("numTargetRowsUpdated"), + metricValue("numTargetRowsInserted"), + metricValue("numTargetRowsMatchedUpdated"), + metricValue("numTargetRowsMatchedDeleted"), + metricValue("numTargetRowsNotMatchedBySourceUpdated"), + metricValue("numTargetRowsNotMatchedBySourceDeleted"))) + case None => + batchWrite.commit(messages) + } } } } diff --git a/spark/src/main/spark-4.1/org/apache/comet/iceberg/IcebergDeltaLogicalFieldsShim.scala b/spark/src/main/spark-4.1/org/apache/comet/iceberg/IcebergDeltaLogicalFieldsShim.scala new file mode 100644 index 00000000000..8aab1f72b57 --- /dev/null +++ b/spark/src/main/spark-4.1/org/apache/comet/iceberg/IcebergDeltaLogicalFieldsShim.scala @@ -0,0 +1,47 @@ +/* + * 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.iceberg + +import org.apache.spark.sql.catalyst.analysis.NamedRelation +import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, WriteDelta} +import org.apache.spark.sql.connector.write.RowLevelOperation.Command.{DELETE, MERGE, UPDATE} + +/** Spark 4.1 exposes the WriteDelta fields directly; keep extraction fully typed. */ +private[iceberg] object IcebergDeltaLogicalFieldsShim extends IcebergDeltaLogicalFieldsShimApi { + override def extract(plan: LogicalPlan): Option[DeltaLogicalFields] = plan match { + case delta: WriteDelta => + val command = delta.operation.command match { + case DELETE => Some(DeltaDelete) + case UPDATE => Some(DeltaUpdate) + case MERGE => Some(DeltaMerge) + case _ => None + } + command.map { value => + val relation: NamedRelation = delta.originalTable + DeltaLogicalFields( + query = delta.query, + originalTable = relation, + projections = delta.projections, + write = delta.write, + command = Some(value)) + } + case _ => None + } +} diff --git a/spark/src/main/spark-4.1/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala b/spark/src/main/spark-4.1/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala new file mode 100644 index 00000000000..d4d145fc4ad --- /dev/null +++ b/spark/src/main/spark-4.1/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala @@ -0,0 +1,30 @@ +/* + * 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.iceberg + +import org.apache.spark.SparkContext +import org.apache.spark.sql.execution.metric.SQLMetric + +/** Retry-aware row-level metrics are added in Spark 4.2. */ +object IcebergSemanticMetricsShim { + def deltaMetrics(sc: SparkContext, command: Option[DeltaCommand]): Map[String, SQLMetric] = + Map.empty + def value(metric: SQLMetric): Long = metric.value +} diff --git a/spark/src/main/spark-4.1/org/apache/spark/sql/comet/IcebergDeltaWriteSummaryShim.scala b/spark/src/main/spark-4.1/org/apache/spark/sql/comet/IcebergDeltaWriteSummaryShim.scala new file mode 100644 index 00000000000..36d70e65670 --- /dev/null +++ b/spark/src/main/spark-4.1/org/apache/spark/sql/comet/IcebergDeltaWriteSummaryShim.scala @@ -0,0 +1,34 @@ +/* + * 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.connector.write.{BatchWrite, WriterCommitMessage} +import org.apache.spark.sql.execution.SparkPlan + +import org.apache.comet.iceberg.DeltaCommand + +/** Spark 4.1 has MERGE summaries; UPDATE and DELETE summaries arrive with Spark 4.2. */ +private[comet] object IcebergDeltaWriteSummaryShim { + def commit( + batchWrite: BatchWrite, + messages: Array[WriterCommitMessage], + query: SparkPlan, + command: Option[DeltaCommand]): Boolean = false +} diff --git a/spark/src/main/spark-4.2/org/apache/comet/iceberg/IcebergDeltaLogicalFieldsShim.scala b/spark/src/main/spark-4.2/org/apache/comet/iceberg/IcebergDeltaLogicalFieldsShim.scala new file mode 100644 index 00000000000..8964df79ab2 --- /dev/null +++ b/spark/src/main/spark-4.2/org/apache/comet/iceberg/IcebergDeltaLogicalFieldsShim.scala @@ -0,0 +1,47 @@ +/* + * 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.iceberg + +import org.apache.spark.sql.catalyst.analysis.NamedRelation +import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, WriteDelta} +import org.apache.spark.sql.connector.write.RowLevelOperation.Command.{DELETE, MERGE, UPDATE} + +/** Spark 4.2 exposes the row-level fields directly; keep this extraction fully typed. */ +private[iceberg] object IcebergDeltaLogicalFieldsShim extends IcebergDeltaLogicalFieldsShimApi { + override def extract(plan: LogicalPlan): Option[DeltaLogicalFields] = plan match { + case delta: WriteDelta => + val command = delta.operation.command match { + case DELETE => Some(DeltaDelete) + case UPDATE => Some(DeltaUpdate) + case MERGE => Some(DeltaMerge) + case _ => None + } + command.map { value => + val relation: NamedRelation = delta.originalTable + DeltaLogicalFields( + query = delta.query, + originalTable = relation, + projections = delta.projections, + write = delta.write, + command = Some(value)) + } + case _ => None + } +} diff --git a/spark/src/main/spark-4.2/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala b/spark/src/main/spark-4.2/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala new file mode 100644 index 00000000000..4808c4ebfb2 --- /dev/null +++ b/spark/src/main/spark-4.2/org/apache/comet/iceberg/IcebergSemanticMetricsShim.scala @@ -0,0 +1,45 @@ +/* + * 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.iceberg + +import org.apache.spark.SparkContext +import org.apache.spark.sql.execution.metric.{SQLLastAttemptMetric, SQLLastAttemptMetrics, SQLMetric} + +/** Spark 4.2 uses last-attempt accumulators for row-level write summaries. */ +object IcebergSemanticMetricsShim { + def deltaMetrics(sc: SparkContext, command: Option[DeltaCommand]): Map[String, SQLMetric] = + command match { + case Some(DeltaUpdate) => + Map( + "numUpdatedRows" -> SQLLastAttemptMetrics.createMetric(sc, "number of updated rows"), + "numCopiedRows" -> SQLLastAttemptMetrics.createMetric(sc, "number of copied rows")) + case Some(DeltaDelete) => + Map( + "numDeletedRows" -> SQLLastAttemptMetrics.createMetric(sc, "number of deleted rows"), + "numCopiedRows" -> SQLLastAttemptMetrics.createMetric(sc, "number of copied rows")) + case _ => Map.empty + } + + def value(metric: SQLMetric): Long = metric match { + case lastAttempt: SQLLastAttemptMetric => + lastAttempt.lastAttemptValueForHighestRDDId().getOrElse(lastAttempt.value) + case other => other.value + } +} diff --git a/spark/src/main/spark-4.2/org/apache/spark/sql/comet/IcebergDeltaWriteSummaryShim.scala b/spark/src/main/spark-4.2/org/apache/spark/sql/comet/IcebergDeltaWriteSummaryShim.scala new file mode 100644 index 00000000000..72879e35ebc --- /dev/null +++ b/spark/src/main/spark-4.2/org/apache/spark/sql/comet/IcebergDeltaWriteSummaryShim.scala @@ -0,0 +1,58 @@ +/* + * 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.connector.write.{BatchWrite, DeleteSummaryImpl, UpdateSummaryImpl, WriterCommitMessage} +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.metric.SQLMetric + +import org.apache.comet.iceberg.{DeltaCommand, DeltaDelete, DeltaUpdate, IcebergSemanticMetricsShim} + +/** Spark 4.2 DELETE/UPDATE summary handling for the JVM PositionDelta writer. */ +private[comet] object IcebergDeltaWriteSummaryShim extends AdaptiveSparkPlanHelper { + def commit( + batchWrite: BatchWrite, + messages: Array[WriterCommitMessage], + query: SparkPlan, + command: Option[DeltaCommand]): Boolean = command match { + case Some(DeltaUpdate) => + deltaWriterMetrics(query).exists { metrics => + batchWrite.commit( + messages, + UpdateSummaryImpl(value(metrics, "numUpdatedRows"), value(metrics, "numCopiedRows"))) + true + } + case Some(DeltaDelete) => + deltaWriterMetrics(query).exists { metrics => + batchWrite.commit( + messages, + DeleteSummaryImpl(value(metrics, "numDeletedRows"), value(metrics, "numCopiedRows"))) + true + } + case _ => false + } + + private def deltaWriterMetrics(query: SparkPlan): Option[Map[String, SQLMetric]] = + collectFirst(query) { case writer: IcebergWriteExec => writer.metrics } + + private def value(metrics: Map[String, SQLMetric], name: String): Long = + metrics.get(name).map(IcebergSemanticMetricsShim.value).getOrElse(-1L) +} diff --git a/spark/src/main/spark-4.2/org/apache/spark/sql/comet/IcebergTransactionalCommitExec.scala b/spark/src/main/spark-4.2/org/apache/spark/sql/comet/IcebergTransactionalCommitExec.scala index 7dec0be1d23..a749217393d 100644 --- a/spark/src/main/spark-4.2/org/apache/spark/sql/comet/IcebergTransactionalCommitExec.scala +++ b/spark/src/main/spark-4.2/org/apache/spark/sql/comet/IcebergTransactionalCommitExec.scala @@ -33,7 +33,8 @@ private[org] final class IcebergTransactionalCommitExec private ( delegate.batchWrite, delegate.write, delegate.refreshCache, - delegate.child) + delegate.child, + delegate.command) with TransactionalExec { override def withTransaction(txn: Option[Transaction]): SparkPlan = copyWith(transaction = txn) @@ -62,13 +63,14 @@ private[org] final class IcebergTransactionalCommitExec private ( } override def makeCopy(newArgs: Array[AnyRef]): IcebergTransactionalCommitExec = { - require(newArgs.length == 4, s"expected 4 IcebergCommitExec arguments, got ${newArgs.length}") + require(newArgs.length == 5, s"expected 5 IcebergCommitExec arguments, got ${newArgs.length}") val copied = new IcebergTransactionalCommitExec( IcebergCommitExec( newArgs(0).asInstanceOf[org.apache.spark.sql.connector.write.BatchWrite], newArgs(1).asInstanceOf[org.apache.spark.sql.connector.write.Write], newArgs(2).asInstanceOf[IcebergCommitExec.RefreshCache], - newArgs(3).asInstanceOf[SparkPlan]), + newArgs(3).asInstanceOf[SparkPlan], + newArgs(4).asInstanceOf[Option[org.apache.comet.iceberg.DeltaCommand]]), transaction) copied.copyTagsFrom(this) copied @@ -85,7 +87,7 @@ private[org] final class IcebergTransactionalCommitExec private ( child: SparkPlan = this.child, transaction: Option[Transaction] = this.transaction): IcebergTransactionalCommitExec = new IcebergTransactionalCommitExec( - IcebergCommitExec(batchWrite, write, refreshCache, child), + IcebergCommitExec(batchWrite, write, refreshCache, child, command), transaction) } diff --git a/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergDeltaLogicalShim.scala b/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergDeltaLogicalShim.scala new file mode 100644 index 00000000000..43354cbec61 --- /dev/null +++ b/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergDeltaLogicalShim.scala @@ -0,0 +1,34 @@ +/* + * 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.iceberg + +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan + +/** Spark 4.x uses the version-neutral WriteDelta descriptor for Iceberg row-level writes. */ +private[iceberg] object IcebergDeltaLogicalShim extends IcebergDeltaLogicalShimApi { + private val writeDeltaClass = "org.apache.spark.sql.catalyst.plans.logical.WriteDelta" + + override def extract(plan: LogicalPlan): Option[DeltaLogicalFields] = + if (plan.getClass.getName == writeDeltaClass) { + IcebergDeltaLogicalFieldsShim.extract(plan) + } else { + None + } +} diff --git a/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergDeltaWriterShim.scala b/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergDeltaWriterShim.scala new file mode 100644 index 00000000000..ceb6ff18460 --- /dev/null +++ b/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergDeltaWriterShim.scala @@ -0,0 +1,70 @@ +/* + * 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.iceberg + +import org.apache.spark.SparkException +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.util.RowDeltaUtils.{DELETE_OPERATION, INSERT_OPERATION, REINSERT_OPERATION, UPDATE_OPERATION} +import org.apache.spark.sql.connector.write.DeltaWriter + +/** Spark 4.x adds REINSERT to the WriteDelta operation stream. */ +object IcebergDeltaWriterShim { + val OperationCodes: DeltaOperationCodes = + DeltaOperationCodes( + delete = DELETE_OPERATION, + update = UPDATE_OPERATION, + insert = INSERT_OPERATION, + reinsert = Some(REINSERT_OPERATION)) + + def writeOperation( + writer: DeltaWriter[InternalRow], + operation: Int, + row: InternalRow, + info: WriteDeltaDispatchInfo): Unit = { + def metadata(): InternalRow = info.metadataProjection.map { projection => + projection.project(row) + projection + }.orNull + + operation match { + case code if code == info.operationCodes.delete => + info.rowIdProjection.project(row) + writer.delete(metadata(), info.rowIdProjection) + case code if code == info.operationCodes.update => + val data = info.rowProjection.getOrElse( + throw new IllegalArgumentException("UPDATE operation has no row projection")) + data.project(row) + info.rowIdProjection.project(row) + writer.update(metadata(), info.rowIdProjection, data) + case code if code == info.operationCodes.insert => + val data = info.rowProjection.getOrElse( + throw new IllegalArgumentException("INSERT operation has no row projection")) + data.project(row) + writer.insert(data) + case code if info.operationCodes.reinsert.contains(code) => + val data = info.rowProjection.getOrElse( + throw new IllegalArgumentException("REINSERT operation has no row projection")) + data.project(row) + writer.reinsert(metadata(), data) + case _ => + throw new SparkException(s"Unexpected operation ID: $operation") + } + } +} diff --git a/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala b/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala index 67a4363870d..5e5641fbb2e 100644 --- a/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala +++ b/spark/src/main/spark-4.x/org/apache/spark/sql/comet/shims/ShimSparkErrorConverter.scala @@ -318,6 +318,12 @@ trait ShimSparkErrorConverter { params("groupCount").toString.toInt, params("groupIndex").toString.toInt)) + case "InvalidUrl" => + Some( + QueryExecutionErrors.invalidUrlError( + UTF8String.fromString(params("url").toString), + new java.net.URISyntaxException(params("url").toString, "Invalid URL"))) + case "DatatypeCannotOrder" => Some( QueryExecutionErrors.orderedOperationUnsupportedByDataTypeError( diff --git a/spark/src/test/java/org/apache/comet/hadoop/fs/FakeWasbSchemeFileSystem.java b/spark/src/test/java/org/apache/comet/hadoop/fs/FakeWasbSchemeFileSystem.java new file mode 100644 index 00000000000..95ff00b818a --- /dev/null +++ b/spark/src/test/java/org/apache/comet/hadoop/fs/FakeWasbSchemeFileSystem.java @@ -0,0 +1,51 @@ +/* + * 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.hadoop.fs; + +import java.net.URI; + +import org.apache.hadoop.fs.RawLocalFileSystem; + +/** + * A local-disk-backed FileSystem that reports the {@code wasb} scheme, so a test can create and + * read an Iceberg table under a {@code wasb://} warehouse without Azure. Used to assert that the + * Iceberg scan and write gates decline a table whose storage scheme the native Iceberg storage + * factory has no arm for, instead of claiming it and failing at execution. + */ +public class FakeWasbSchemeFileSystem extends RawLocalFileSystem { + + public static final String PREFIX = "wasb://fake-container"; + + public FakeWasbSchemeFileSystem() { + // Avoid `URI scheme is not "file"` error on + // RawLocalFileSystem$DeprecatedRawLocalFileStatus.getOwner + RawLocalFileSystem.useStatIfAvailable(); + } + + @Override + public String getScheme() { + return "wasb"; + } + + @Override + public URI getUri() { + return URI.create(PREFIX); + } +} diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_contains.sql b/spark/src/test/resources/sql-tests/expressions/array/array_contains.sql index e0fb8709b93..6a38f7efdd1 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/array_contains.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/array_contains.sql @@ -59,8 +59,8 @@ SELECT array_contains(cast(NULL as array), val) FROM test_array_contains -- ============================================================ -- Floating-point elements: Spark compares with ordering.equiv, so -0.0 == +0.0 and NaN == NaN. --- Native array_contains compares raw Arrow values bitwise, so float/double element types route --- through the JVM codegen dispatcher (Spark's own doGenCode) to stay native and Spark-exact. +-- Flat float/double arrays run Comet's spark_array_contains, which compares the same way; nested +-- float elements route through the JVM codegen dispatcher (Spark's own doGenCode). -- ============================================================ statement @@ -87,6 +87,6 @@ SELECT array_contains(array(-0.0D, 1.0D), 0.0D), array_contains(array(cast('-0.0' as float), 1.0F), cast('0.0' as float)), array_contains(array(cast('NaN' as float), 1.0F), cast('NaN' as float)) --- nested float element: the element type is array, so the float decline must recurse +-- nested float element: the element type is array, so the dispatch decision must recurse query SELECT array_contains(array(array(0.0D), array(1.0D)), array(-0.0D)) diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_contains_collation.sql b/spark/src/test/resources/sql-tests/expressions/array/array_contains_collation.sql new file mode 100644 index 00000000000..0a4d4306bc1 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_contains_collation.sql @@ -0,0 +1,32 @@ +-- 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. + +-- A collated string field inside a struct element compares by its collation in Spark. Nested +-- elements stay on the codegen dispatcher, which keeps that. Collation syntax requires Spark 4.0+. +-- MinSparkVersion: 4.0 + +statement +CREATE TABLE ac_collation(d DOUBLE, s STRING) USING parquet + +statement +INSERT INTO ac_collation VALUES (0.0D, 'a'), (1.0D, 'b'), (NULL, NULL) + +query expect_dispatch(array_contains) +SELECT array_contains( + array(named_struct('d', d, 's', CAST(s AS STRING COLLATE UTF8_LCASE))), + named_struct('d', d, 's', CAST(upper(s) AS STRING COLLATE UTF8_LCASE))) +FROM ac_collation diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_contains_eager_eval.sql b/spark/src/test/resources/sql-tests/expressions/array/array_contains_eager_eval.sql new file mode 100644 index 00000000000..359f4b83803 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_contains_eager_eval.sql @@ -0,0 +1,34 @@ +-- 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. + +-- Spark skips array_contains' value when the array is null, so a value that throws is not +-- evaluated for that row. The native kernel evaluates both arguments first, so a value other than +-- a literal or a column read stays on the codegen dispatcher when the array can be null. +-- Config: spark.sql.ansi.enabled=true + +statement +CREATE TABLE ac_eager(a ARRAY, s STRING) USING parquet + +statement +INSERT INTO ac_eager VALUES (NULL, 'bad'), (array(1.0D), '1') + +query expect_dispatch(array_contains) +SELECT array_contains(a, CAST(s AS DOUBLE)) FROM ac_eager + +-- A column read is safe to evaluate early, so it stays native +query expect_native(array_contains) +SELECT array_contains(a, 1.0D), array_contains(a, CAST(NULL AS DOUBLE)) FROM ac_eager diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_contains_floating_point.sql b/spark/src/test/resources/sql-tests/expressions/array/array_contains_floating_point.sql new file mode 100644 index 00000000000..010608cecfb --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_contains_floating_point.sql @@ -0,0 +1,85 @@ +-- 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. + +-- array_contains compares floats the way Spark's generated code does: -0.0 equals 0.0, and every +-- NaN equals every other NaN, at any depth of an array or struct element. Flat FLOAT and DOUBLE +-- arrays run natively; nested float elements go through the codegen dispatcher. The result is true when +-- an element matches, null when none does and the array holds a null element, and false +-- otherwise; a null array or value gives null. +-- +-- `-d` flips the sign bit, so `-d` of the NaN row is a NaN with the sign bit set on every +-- platform. That is the NaN that arithmetic produces on x86-64. + +statement +CREATE TABLE ac_float(id INT, d DOUBLE, f FLOAT) USING parquet + +statement +INSERT INTO ac_float VALUES + (1, 0.0D, float('0.0')), + (2, double('-0.0'), float('-0.0')), + (3, double('NaN'), float('NaN')), + (4, 1.0D, float('1.0')), + (5, NULL, NULL) + +-- A constant value: either zero finds both zeros, and a NaN finds both NaNs +query expect_native(array_contains) +SELECT id, array_contains(array(d, -d), 0.0D), array_contains(array(d, -d), -0.0D), + array_contains(array(-d, 1.0D), double('NaN')), array_contains(array(d, 1.0D), 2.0D) +FROM ac_float + +query expect_native(array_contains) +SELECT id, array_contains(array(f, -f), 0.0F), array_contains(array(-f, 2.0F), float('NaN')) +FROM ac_float + +-- A value per row +query expect_native(array_contains) +SELECT id, array_contains(array(0.0D, 1.0D), d), array_contains(array(double('NaN'), 1.0D), -d), + array_contains(array(float('-0.0'), 2.0F), f) +FROM ac_float + +-- Null elements: true when another element matches, null when nothing matches +query expect_native(array_contains) +SELECT id, array_contains(array(d, NULL), 0.0D), array_contains(array(2.0D, NULL), d), + array_contains(array(NULL, -d), double('NaN')) +FROM ac_float + +-- A null array or a null value gives null; an empty array gives false +query expect_native(array_contains) +SELECT id, array_contains(IF(d IS NULL, NULL, array(d, -d)), 0.0D), + array_contains(array(d, -d), CAST(NULL AS DOUBLE)), array_contains(array(), d) +FROM ac_float + +-- Arrays as elements compare their floats the same way. Nested float elements stay on the codegen +-- dispatcher. +query expect_dispatch(array_contains) +SELECT id, array_contains(array(array(-d), array(1.0D)), array(d)), + array_contains(array(array(d, 1.0D), array(2.0D)), array(-d, 1.0D)), + array_contains(array(array(d), NULL), array(2.0D)) +FROM ac_float + +-- Struct elements too +query expect_dispatch(array_contains) +SELECT id, array_contains(array(named_struct('x', -d, 'y', 1), named_struct('x', 2.0D, 'y', 1)), + named_struct('x', d, 'y', 1)) +FROM ac_float + +-- Only literals, which the suite does not constant fold. `-0.0D` is a literal with the sign bit +-- set, while `double('NaN')` is a cast. +query expect_native(array_contains) +SELECT array_contains(array(0.0D, 1.0D), -0.0D), array_contains(array(-0.0D), 0.0D), + array_contains(array(double('NaN'), 1.0D), double('NaN')), + array_contains(array(1.0D, NULL), -0.0D) diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_element_equality_collation.sql b/spark/src/test/resources/sql-tests/expressions/array/array_element_equality_collation.sql new file mode 100644 index 00000000000..e8073461829 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_element_equality_collation.sql @@ -0,0 +1,90 @@ +-- 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. + +-- MinSparkVersion: 4.0 +-- Config: spark.comet.exec.scalaUDF.codegen.enabled=true +-- Config: spark.comet.expression.ArrayContains.allowIncompatible=false +-- Config: spark.comet.expression.ArraysOverlap.allowIncompatible=false +-- Config: spark.comet.expression.ArrayDistinct.allowIncompatible=false +-- Config: spark.comet.expression.ArrayUnion.allowIncompatible=false + +-- array_contains, arrays_overlap, array_distinct and array_union compare array elements for +-- equality. The native kernels compare strings by raw bytes, so under a non-UTF8_BINARY collation +-- they would miss 'a' = 'A' (UTF8_LCASE) or 'x ' = 'x' (UTF8_BINARY_RTRIM). Collated +-- array_contains and arrays_overlap route through the JVM codegen dispatcher, which runs Spark's +-- collation-aware comparison. array_distinct and array_union return arrays, where dispatch costs +-- more than a projection fallback (see ArraySetSupport), so they fall back to Spark instead. + +statement +CREATE TABLE test_array_eq_collation(id int, a string, b string) USING parquet + +statement +INSERT INTO test_array_eq_collation VALUES + (1, 'a', 'A'), + (2, 'x ', 'x'), + (3, 'b', 'c'), + (4, NULL, 'a'), + (5, 'A', NULL) + +-- UTF8_BINARY strings stay on the native kernels. +query expect_native(array_contains,arrays_overlap,array_distinct,array_union) +SELECT id, array_contains(array(a), b), arrays_overlap(array(a), array(b)), + array_distinct(array(a, b)), array_union(array(a), array(b)) +FROM test_array_eq_collation + +-- UTF8_LCASE: 'a' and 'A' are equal. +query expect_dispatch(array_contains) +SELECT id, array_contains(array(CAST(a AS STRING COLLATE UTF8_LCASE)), CAST(b AS STRING COLLATE UTF8_LCASE)) +FROM test_array_eq_collation + +query expect_dispatch(arrays_overlap) +SELECT id, arrays_overlap(array(CAST(a AS STRING COLLATE UTF8_LCASE)), array(CAST(b AS STRING COLLATE UTF8_LCASE))) +FROM test_array_eq_collation + +query expect_fallback(native array_distinct compares raw bytes) +SELECT id, array_distinct(array(CAST(a AS STRING COLLATE UTF8_LCASE), CAST(b AS STRING COLLATE UTF8_LCASE))) +FROM test_array_eq_collation + +query expect_fallback(native array_union compares raw bytes) +SELECT id, array_union(array(CAST(a AS STRING COLLATE UTF8_LCASE)), array(CAST(b AS STRING COLLATE UTF8_LCASE))) +FROM test_array_eq_collation + +-- UTF8_BINARY_RTRIM: 'x ' and 'x' are equal. +query expect_dispatch(array_contains) +SELECT id, array_contains(array(CAST(a AS STRING COLLATE UTF8_BINARY_RTRIM)), CAST(b AS STRING COLLATE UTF8_BINARY_RTRIM)) +FROM test_array_eq_collation + +query expect_dispatch(arrays_overlap) +SELECT id, arrays_overlap(array(CAST(a AS STRING COLLATE UTF8_BINARY_RTRIM)), array(CAST(b AS STRING COLLATE UTF8_BINARY_RTRIM))) +FROM test_array_eq_collation + +query expect_fallback(native array_distinct compares raw bytes) +SELECT id, array_distinct(array(CAST(a AS STRING COLLATE UTF8_BINARY_RTRIM), CAST(b AS STRING COLLATE UTF8_BINARY_RTRIM))) +FROM test_array_eq_collation + +query expect_fallback(native array_union compares raw bytes) +SELECT id, array_union(array(CAST(a AS STRING COLLATE UTF8_BINARY_RTRIM)), array(CAST(b AS STRING COLLATE UTF8_BINARY_RTRIM))) +FROM test_array_eq_collation + +-- Collation detection also recurses into struct fields of the array elements. +query expect_dispatch(array_contains,arrays_overlap) +SELECT id, + array_contains(array(named_struct('s', CAST(a AS STRING COLLATE UTF8_LCASE))), + named_struct('s', CAST(b AS STRING COLLATE UTF8_LCASE))), + arrays_overlap(array(named_struct('s', CAST(a AS STRING COLLATE UTF8_LCASE))), + array(named_struct('s', CAST(b AS STRING COLLATE UTF8_LCASE)))) +FROM test_array_eq_collation diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_element_equality_collation_allow_incompatible.sql b/spark/src/test/resources/sql-tests/expressions/array/array_element_equality_collation_allow_incompatible.sql new file mode 100644 index 00000000000..7bedec72ce9 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_element_equality_collation_allow_incompatible.sql @@ -0,0 +1,45 @@ +-- 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. + +-- MinSparkVersion: 4.0 +-- Config: spark.comet.exec.scalaUDF.codegen.enabled=true +-- Config: spark.comet.expression.ArrayContains.allowIncompatible=true +-- Config: spark.comet.expression.ArraysOverlap.allowIncompatible=true +-- Config: spark.comet.expression.ArrayDistinct.allowIncompatible=true +-- Config: spark.comet.expression.ArrayUnion.allowIncompatible=true + +-- allowIncompatible=true opts collated inputs into the native kernels, which compare raw bytes. +-- The data only holds values whose bytewise and collation-aware equality agree, so the answers +-- still match Spark; the point is that the native path is taken. + +statement +CREATE TABLE test_array_eq_collation(id int, a string, b string) USING parquet + +statement +INSERT INTO test_array_eq_collation VALUES (1, 'a', 'a'), (2, 'b', 'c'), (3, NULL, 'a') + +query expect_native(array_contains,arrays_overlap) +SELECT id, + array_contains(array(CAST(a AS STRING COLLATE UTF8_LCASE)), CAST(b AS STRING COLLATE UTF8_LCASE)), + arrays_overlap(array(CAST(a AS STRING COLLATE UTF8_LCASE)), array(CAST(b AS STRING COLLATE UTF8_LCASE))) +FROM test_array_eq_collation + +query expect_native(array_distinct,array_union) +SELECT id, + array_distinct(array(CAST(a AS STRING COLLATE UTF8_LCASE), CAST(b AS STRING COLLATE UTF8_LCASE))), + array_union(array(CAST(a AS STRING COLLATE UTF8_LCASE)), array(CAST(b AS STRING COLLATE UTF8_LCASE))) +FROM test_array_eq_collation diff --git a/spark/src/test/resources/sql-tests/expressions/array/array_element_equality_collation_disabled.sql b/spark/src/test/resources/sql-tests/expressions/array/array_element_equality_collation_disabled.sql new file mode 100644 index 00000000000..e07ad2f8a2c --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/array_element_equality_collation_disabled.sql @@ -0,0 +1,49 @@ +-- 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. + +-- MinSparkVersion: 4.0 +-- Config: spark.comet.exec.scalaUDF.codegen.enabled=false +-- Config: spark.comet.expression.ArrayContains.allowIncompatible=false +-- Config: spark.comet.expression.ArraysOverlap.allowIncompatible=false +-- Config: spark.comet.expression.ArrayDistinct.allowIncompatible=false +-- Config: spark.comet.expression.ArrayUnion.allowIncompatible=false + +-- With the JVM codegen dispatcher disabled, collated array_contains and arrays_overlap have no +-- Spark-compatible Comet path and fall back to Spark. array_distinct and array_union fall back +-- whether or not the dispatcher is enabled. + +statement +CREATE TABLE test_array_eq_collation(id int, a string, b string) USING parquet + +statement +INSERT INTO test_array_eq_collation VALUES (1, 'a', 'A'), (2, 'x ', 'x'), (3, 'b', 'c'), (4, NULL, 'a') + +query expect_fallback(array_contains: spark.comet.exec.scalaUDF.codegen.enabled=false) +SELECT id, array_contains(array(CAST(a AS STRING COLLATE UTF8_LCASE)), CAST(b AS STRING COLLATE UTF8_LCASE)) +FROM test_array_eq_collation + +query expect_fallback(arrays_overlap: spark.comet.exec.scalaUDF.codegen.enabled=false) +SELECT id, arrays_overlap(array(CAST(a AS STRING COLLATE UTF8_BINARY_RTRIM)), array(CAST(b AS STRING COLLATE UTF8_BINARY_RTRIM))) +FROM test_array_eq_collation + +query expect_fallback(native array_distinct compares raw bytes) +SELECT id, array_distinct(array(CAST(a AS STRING COLLATE UTF8_LCASE), CAST(b AS STRING COLLATE UTF8_LCASE))) +FROM test_array_eq_collation + +query expect_fallback(native array_union compares raw bytes) +SELECT id, array_union(array(CAST(a AS STRING COLLATE UTF8_BINARY_RTRIM)), array(CAST(b AS STRING COLLATE UTF8_BINARY_RTRIM))) +FROM test_array_eq_collation diff --git a/spark/src/test/resources/sql-tests/expressions/array/explode.sql b/spark/src/test/resources/sql-tests/expressions/array/explode.sql index b1bdb8b7493..881186e264e 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/explode.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/explode.sql @@ -328,7 +328,7 @@ CREATE TABLE test_explode_empty(id int, arr array) USING parquet query SELECT id, explode_outer(arr) FROM test_explode_empty --- ===== Map falls back to Spark: outer form must also fall back ===== +-- ===== Map inputs: outer form preserves empty and null maps ===== statement CREATE TABLE test_explode_map(id int, m map) USING parquet @@ -339,7 +339,7 @@ INSERT INTO test_explode_map VALUES (2, map()), (3, NULL) -query expect_fallback(Comet only supports explode/explode_outer for arrays, not maps) +query SELECT id, explode_outer(m) FROM test_explode_map -- ===== TIMESTAMP_NTZ (Spark 3.4+): distinct Arrow layout (no tz) from LTZ ===== diff --git a/spark/src/test/resources/sql-tests/expressions/array/folded_array_map_literals.sql b/spark/src/test/resources/sql-tests/expressions/array/folded_array_map_literals.sql new file mode 100644 index 00000000000..f63cee958d1 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/folded_array_map_literals.sql @@ -0,0 +1,33 @@ +-- 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. + +-- ConstantFolding: enabled + +statement +CREATE TABLE folded_array_map_literals(id INT, value BIGINT) USING parquet + +statement +INSERT INTO folded_array_map_literals SELECT CAST(id AS INT), id FROM range(0, 3, 1, 1) + +-- The whole array folds to one ArrayType(MapType(...), containsNull=true) literal. +-- The NULL element and the rebuilt populated map must retain identical map types. +-- Three rows in one file exercise broadcasting the folded literal across a batch. +query +SELECT array(map(1, 2), NULL) AS arr FROM folded_array_map_literals + +query +SELECT array(NULL, map(1, 2)) AS arr FROM folded_array_map_literals diff --git a/spark/src/test/resources/sql-tests/expressions/array/posexplode.sql b/spark/src/test/resources/sql-tests/expressions/array/posexplode.sql index eb537586f0e..d31c00619c4 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/posexplode.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/posexplode.sql @@ -93,8 +93,8 @@ INSERT INTO test_posexplode_map VALUES (1, map('a', 1, 'b', 2)), (2, map('c', 3)) --- posexplode over a map falls back to Spark (Comet only supports array inputs, not maps) -query expect_fallback(Comet only supports explode/explode_outer for arrays, not maps) +-- Map entries expand to position, key and value columns. +query SELECT id, posexplode(m) FROM test_posexplode_map -- ===== posexplode_outer across non-int element types ===== diff --git a/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_disabled.sql b/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_disabled.sql index a6e9499c047..b3991c1a1dc 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_disabled.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_disabled.sql @@ -44,7 +44,8 @@ SELECT array_join(words, '-') FROM routing_arrays query expect_native(reverse) SELECT reverse(a) FROM routing_arrays -query expect_fallback(array_contains: spark.comet.exec.scalaUDF.codegen.enabled=false) +-- array_contains on floats follows Spark natively, so it needs no dispatcher. +query expect_native(array_contains) SELECT array_contains(f, 1.0D) FROM routing_arrays -- sort_array on floats follows Spark natively, even in strict floating-point mode. diff --git a/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_enabled.sql b/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_enabled.sql index 1e688857a9e..c401a040a9c 100644 --- a/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_enabled.sql +++ b/spark/src/test/resources/sql-tests/expressions/array/routing_arrays_enabled.sql @@ -44,7 +44,8 @@ SELECT array_join(words, '-') FROM routing_arrays query expect_native(reverse) SELECT reverse(a) FROM routing_arrays -query expect_dispatch(array_contains) +-- array_contains on floats follows Spark natively, even in strict floating-point mode. +query expect_native(array_contains) SELECT array_contains(f, 1.0D) FROM routing_arrays -- sort_array on floats follows Spark natively, even in strict floating-point mode. diff --git a/spark/src/test/resources/sql-tests/expressions/bitwise/shift_mixed_types.sql b/spark/src/test/resources/sql-tests/expressions/bitwise/shift_mixed_types.sql new file mode 100644 index 00000000000..a8100c48392 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/bitwise/shift_mixed_types.sql @@ -0,0 +1,46 @@ +-- 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. + +-- ConfigMatrix: parquet.enable.dictionary=false,true + +-- BIGINT values with INT amounts exercise Comet's cast of the shift amount. +-- Migrated from CometBitwiseExpressionSuite's mixed-type shift test (#6634). +statement +CREATE TABLE test_shift_mixed_types(v bigint, amount int) USING parquet + +statement +INSERT INTO test_shift_mixed_types +SELECT /*+ COALESCE(1) */ v, amount +FROM VALUES + (1111, 2), + (1111, 2), + (3333, 4), + (5555, 6), + (4294967296, 1), + (-4294967296, 1), + (9223372036854775807, 1), + (-9223372036854775808, 1), + (0, 0), + (NULL, 2), + (1111, NULL) +AS input(v, amount) + +query +SELECT shiftright(v, 2), shiftright(v, amount) FROM test_shift_mixed_types + +query +SELECT shiftleft(v, 2), shiftleft(v, amount) FROM test_shift_mixed_types diff --git a/spark/src/test/resources/sql-tests/expressions/conditional/float_comparisons.sql b/spark/src/test/resources/sql-tests/expressions/conditional/float_comparisons.sql new file mode 100644 index 00000000000..61b4b36f0e9 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/conditional/float_comparisons.sql @@ -0,0 +1,207 @@ +-- 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. + +-- Comparisons follow Spark's SQL ordering for floats (SQLOrderingUtil.compareDoubles): -0.0 +-- equals 0.0, every NaN equals every other NaN, and NaN sorts above every other value. This +-- must hold in every operator that evaluates a comparison, not only in Project and Filter. +-- +-- `-d` flips the sign bit, so `-d` of the NaN row is a NaN with the sign bit set on every +-- platform. That is the NaN that arithmetic produces on x86-64, and Arrow's total order sorts +-- it below -Infinity. +-- +-- A scan pushes its filters into the Parquet reader, which with row-level pushdown drops rows +-- itself, before Spark's Filter above the scan sees them. +-- ConfigMatrix: spark.comet.parquet.rowFilterPushdown.enabled=false,true + +statement +CREATE TABLE float_cmp(id INT, d DOUBLE, f FLOAT) USING parquet + +statement +INSERT INTO float_cmp VALUES + (1, 0.0D, float('0.0')), + (2, double('-0.0'), float('-0.0')), + (3, double('NaN'), float('NaN')), + (4, 1.0D, float('1.0')), + (5, -1.0D, float('-1.0')), + (6, double('Infinity'), float('Infinity')), + (7, NULL, NULL) + +-- Project. `IS DISTINCT FROM` is `NOT (a <=> b)`, which Comet plans as one native operator. +query +SELECT id, -d = d, -d <=> d, -d < d, -d <= d, -d > d, -d >= d, -d != d, -d IS DISTINCT FROM d +FROM float_cmp + +query +SELECT id, -f = f, -f <=> f, -f < f, -f <= f, -f > f, -f >= f, -f != f, -f IS DISTINCT FROM f +FROM float_cmp + +-- Comparisons against a constant on either side. `-0.0D` and `-0.0F` are literals with the sign +-- bit set, while `double('NaN')` is a cast, because the suite turns off constant folding. +query +SELECT id, -d > 0.0D, -d >= -0.0D, -0.0D = d, -d = double('NaN'), double('NaN') <= -d, + -d < double('-0.0'), -d IS DISTINCT FROM double('NaN'), -0.0D IS DISTINCT FROM d, + -f > 0.0F, -0.0F <=> f, -f = float('NaN'), float('-0.0') >= -f +FROM float_cmp + +-- Filter +query +SELECT id FROM float_cmp WHERE -d > 0.0D + +query +SELECT id FROM float_cmp WHERE -f >= f + +-- A stored column compared with a constant, the filter shape that Parquet pruning reads. The +-- scan still returns the -0.0 and NaN rows that Spark's ordering keeps. +query +SELECT id FROM float_cmp WHERE d >= 0.0D + +query +SELECT id FROM float_cmp WHERE f <= -0.0F + +query +SELECT id FROM float_cmp WHERE d = double('NaN') + +query +SELECT id FROM float_cmp WHERE d > double('Infinity') + +-- Aggregate argument +query +SELECT sum(if(-d > 0.0D, 1, 0)), sum(if(-d = d, 1, 0)), sum(if(-f < f, 1, 0)) FROM float_cmp + +-- Aggregate FILTER clause +query +SELECT count(*) FILTER (WHERE -d >= d), count(*) FILTER (WHERE -f <=> f), + count(*) FILTER (WHERE -d IS DISTINCT FROM d) +FROM float_cmp + +-- Grouping by a comparison +query +SELECT -d > 0.0D AS k, count(*) FROM float_cmp GROUP BY -d > 0.0D + +-- Join conditions, for each join strategy +query +SELECT /*+ BROADCAST(b) */ a.id, b.id FROM float_cmp a JOIN float_cmp b + ON a.id = b.id AND -a.d >= b.d + +query +SELECT /*+ SHUFFLE_HASH(b) */ a.id, b.id FROM float_cmp a JOIN float_cmp b + ON a.id = b.id AND -a.d >= b.d + +query +SELECT /*+ MERGE(b) */ a.id, b.id FROM float_cmp a JOIN float_cmp b + ON a.id = b.id AND -a.f >= b.f + +query +SELECT /*+ BROADCAST(b) */ a.id, b.id FROM float_cmp a JOIN float_cmp b + ON -a.d > b.d WHERE a.id = 3 + +-- Sort key +query +SELECT id FROM float_cmp ORDER BY -d > 0.0D, id + +-- Generate +query +SELECT id, x FROM float_cmp LATERAL VIEW explode(array(-d > 0.0D, -d = d)) v AS x + +-- Window function argument, which Spark evaluates in a Project below the Window +query +SELECT id, sum(if(-d > 0.0D, 1, 0)) OVER (ORDER BY id) FROM float_cmp + +-- Arrays and structs compare their float leaves the same way, for every operator +statement +CREATE TABLE float_cmp_nested( + id INT, + a ARRAY, b ARRAY, + s STRUCT, u STRUCT, + x ARRAY, y ARRAY) USING parquet + +statement +INSERT INTO float_cmp_nested VALUES + (1, array(double('-0.0')), array(0.0D), + named_struct('v', double('-0.0')), named_struct('v', 0.0D), + array(float('-0.0')), array(float('0.0'))), + (2, array(double('NaN')), array(double('Infinity')), + named_struct('v', double('NaN')), named_struct('v', double('Infinity')), + array(float('NaN')), array(float('Infinity'))), + (3, array(1.0D, NULL), array(1.0D, 2.0D), + named_struct('v', NULL), named_struct('v', 1.0D), + array(float('1.0'), NULL), array(float('1.0'), float('2.0'))), + (4, array(1.0D), array(1.0D, double('-0.0')), + named_struct('v', 1.0D), named_struct('v', 1.0D), + array(), array(float('-0.0'))), + (5, NULL, array(0.0D), NULL, named_struct('v', 0.0D), NULL, array(float('0.0'))) + +query +SELECT id, a = b, a <=> b, a < b, a <= b, a > b, a >= b, a IS DISTINCT FROM b +FROM float_cmp_nested + +query +SELECT id, s = u, s <=> u, s < u, s <= u, s > u, s >= u, s IS DISTINCT FROM u +FROM float_cmp_nested + +query +SELECT id, x = y, x <=> y, x < y, x <= y, x > y, x >= y, x IS DISTINCT FROM y +FROM float_cmp_nested + +-- Negating the elements gives the NaN in row 2 the sign bit +query +SELECT id, transform(a, e -> -e) < b, transform(a, e -> -e) <=> a, + transform(a, e -> -e) IS DISTINCT FROM a +FROM float_cmp_nested + +-- Nested comparisons outside Project and Filter +query +SELECT sum(if(a < b, 1, 0)), count(*) FILTER (WHERE s <=> u) FROM float_cmp_nested + +-- Floats two levels deep, in an array of structs and in a struct of arrays +statement +CREATE TABLE float_cmp_deep( + id INT, + a ARRAY>, b ARRAY>, + s STRUCT>, u STRUCT>) USING parquet + +statement +INSERT INTO float_cmp_deep VALUES + (1, array(named_struct('v', double('-0.0'))), array(named_struct('v', 0.0D)), + named_struct('a', array(double('-0.0'))), named_struct('a', array(0.0D))), + (2, array(named_struct('v', double('NaN'))), array(named_struct('v', double('Infinity'))), + named_struct('a', array(double('NaN'))), named_struct('a', array(double('Infinity')))), + (3, array(named_struct('v', 1.0D), named_struct('v', NULL)), + array(named_struct('v', 1.0D), named_struct('v', 2.0D)), + named_struct('a', array(1.0D, NULL)), named_struct('a', array(1.0D, 2.0D))), + (4, array(named_struct('v', 1.0D)), + array(named_struct('v', 1.0D), named_struct('v', double('-0.0'))), + named_struct('a', NULL), named_struct('a', array(double('-0.0')))), + (5, NULL, array(named_struct('v', 0.0D)), NULL, named_struct('a', array(0.0D))) + +query +SELECT id, a = b, a <=> b, a < b, a <= b, a > b, a >= b, a IS DISTINCT FROM b +FROM float_cmp_deep + +query +SELECT id, s = u, s <=> u, s < u, s <= u, s > u, s >= u, s IS DISTINCT FROM u +FROM float_cmp_deep + +-- Negating the leaves gives the NaN in row 2 the sign bit +query +SELECT id, transform(a, e -> named_struct('v', -e.v)) < b, + transform(a, e -> named_struct('v', -e.v)) <=> a, + named_struct('a', transform(s.a, e -> -e)) <=> s +FROM float_cmp_deep + +query +SELECT sum(if(a < b, 1, 0)), count(*) FILTER (WHERE s <=> u) FROM float_cmp_deep diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_date.sql b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_date.sql index 672281627c9..0bc3b431563 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_date.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_date.sql @@ -15,21 +15,77 @@ -- specific language governing permissions and limitations -- under the License. +-- Dictionary-encoded dates reuse the scalar Date32 path for dictionary values. +-- ConfigMatrix: parquet.enable.dictionary=false,true +-- Config: spark.sql.parquet.datetimeRebaseModeInWrite=CORRECTED + statement CREATE TABLE test_trunc_date(d date) USING parquet statement -INSERT INTO test_trunc_date VALUES (date('2024-06-15')), (date('2024-01-01')), (date('2024-12-31')), (NULL) +INSERT INTO test_trunc_date VALUES + (date('2024-05-17')), + -- Monday / Tuesday / Sunday for Monday-based WEEK truncation. + (date('2024-05-13')), + (date('2024-05-14')), + (date('2024-05-19')), + (date('2024-01-01')), + (date('2023-12-31')), + (date('2024-03-31')), + (date('2024-04-01')), + (date('2024-06-30')), + (date('2024-07-01')), + (date('2024-10-01')), + (date('2024-02-29')), + (date('2000-02-29')), + (date('1900-02-28')), + (date('1969-12-31')), + (date('1960-02-29')), + (date('1900-01-01')), + (date('1500-06-15')), + (date('1678-06-01')), + -- Valid Spark Date32 outside TimestampNanosecond's range. + (date('3333-05-17')), + (NULL) + +query +SELECT d, trunc(d, 'YEAR'), trunc(d, 'YYYY'), trunc(d, 'YY') FROM test_trunc_date ORDER BY d + +query +SELECT d, trunc(d, 'QUARTER') FROM test_trunc_date ORDER BY d + +query +SELECT d, trunc(d, 'MONTH'), trunc(d, 'MON'), trunc(d, 'MM') FROM test_trunc_date ORDER BY d + +query +SELECT d, trunc(d, 'WEEK') FROM test_trunc_date ORDER BY d + +query +SELECT d, trunc(d, 'year'), trunc(d, 'Year'), trunc(d, 'yEaR'), trunc(d, 'month'), trunc(d, 'Mon'), trunc(d, 'week') FROM test_trunc_date ORDER BY d + +-- Unsupported scalar formats fall back to Spark's codegen dispatcher and remain NULL-compatible. +query +SELECT d, trunc(d, 'DAY'), trunc(d, 'HOUR'), trunc(d, 'SECOND'), trunc(d, 'invalid'), trunc(d, ' YEAR '), trunc(d, '') FROM test_trunc_date ORDER BY d +-- NULL format is Incompatible on the native path, so this uses the codegen dispatcher. query -SELECT trunc(d, 'year') FROM test_trunc_date +SELECT d, trunc(d, NULL) FROM test_trunc_date ORDER BY d query -SELECT trunc(d, 'month') FROM test_trunc_date +SELECT trunc(NULL, 'YEAR'), trunc(NULL, NULL) +-- Literal arguments exercise the scalar date input branch. query -SELECT trunc(d, 'quarter') FROM test_trunc_date +SELECT + trunc(DATE '2024-05-17', 'YEAR'), + trunc(DATE '2024-05-17', 'QUARTER'), + trunc(DATE '2024-05-17', 'MONTH'), + trunc(DATE '2024-05-17', 'WEEK') --- literal arguments +-- Literal arguments exercise the scalar date input branch. query -SELECT trunc(date('2024-06-15'), 'year'), trunc(date('2024-06-15'), 'month'), trunc(date('2024-06-15'), 'quarter') +SELECT + trunc(DATE '3333-05-17', 'YEAR'), + trunc(DATE '3333-05-17', 'QUARTER'), + trunc(DATE '3333-05-17', 'MONTH'), + trunc(DATE '3333-05-17', 'WEEK') diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp.sql b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp.sql index 1a8bbef9880..662774307ea 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp.sql @@ -15,26 +15,153 @@ -- specific language governing permissions and limitations -- under the License. --- Config: spark.comet.expression.TruncTimestamp.allowIncompatible=true +-- Keep the wide-range fallback fixture independent of far-future JVM/native timezone rules. +-- Config: spark.sql.session.timeZone=UTC +-- Config: spark.sql.parquet.int96RebaseModeInWrite=CORRECTED +-- Config: spark.sql.parquet.datetimeRebaseModeInWrite=CORRECTED +-- Config: spark.sql.parquet.outputTimestampType=TIMESTAMP_MICROS +-- Dictionary-encoded timestamps reuse the scalar timestamp path for dictionary values. +-- ConfigMatrix: parquet.enable.dictionary=false,true statement CREATE TABLE test_trunc_ts(ts timestamp) USING parquet statement -INSERT INTO test_trunc_ts VALUES (timestamp('2024-06-15 10:30:45')), (timestamp('2024-01-01 00:00:00')), (NULL) +INSERT INTO test_trunc_ts VALUES + (timestamp('2024-05-17 12:34:56.123456')), + (timestamp('2024-02-29 23:59:59.999999')), + (timestamp('2000-02-29 00:00:00')), + (timestamp('1900-02-28 00:00:00')), + (timestamp('1969-12-31 23:59:59.123456')), + -- Below the nanosecond lower bound and inside its 370-day truncation margin. + (timestamp('1500-06-15 12:34:56.123456')), + (timestamp('1678-06-01 12:34:56.123456')), + -- Valid Spark timestamp outside TimestampNanosecond's range. + (timestamp('3333-05-17 12:34:56.123456')), + (NULL) query -SELECT date_trunc('year', ts) FROM test_trunc_ts +SELECT ts, date_trunc('YEAR', ts), date_trunc('YYYY', ts), date_trunc('YY', ts) FROM test_trunc_ts ORDER BY ts query -SELECT date_trunc('month', ts) FROM test_trunc_ts +SELECT ts, date_trunc('QUARTER', ts) FROM test_trunc_ts ORDER BY ts query -SELECT date_trunc('day', ts) FROM test_trunc_ts +SELECT ts, date_trunc('MONTH', ts), date_trunc('MON', ts), date_trunc('MM', ts) FROM test_trunc_ts ORDER BY ts query -SELECT date_trunc('hour', ts) FROM test_trunc_ts +SELECT ts, date_trunc('WEEK', ts), date_trunc('DAY', ts), date_trunc('DD', ts) FROM test_trunc_ts ORDER BY ts --- literal arguments query -SELECT date_trunc('year', timestamp('2024-06-15 10:30:45')), date_trunc('month', timestamp('2024-06-15 10:30:45')), date_trunc('day', timestamp('2024-06-15 10:30:45')) +SELECT + ts, + date_trunc('HOUR', ts), + date_trunc('MINUTE', ts), + date_trunc('SECOND', ts), + date_trunc('MILLISECOND', ts), + date_trunc('MICROSECOND', ts) +FROM test_trunc_ts +ORDER BY ts + +query +SELECT + ts, + date_trunc('year', ts), + date_trunc('Year', ts), + date_trunc('yEaR', ts), + date_trunc('month', ts), + date_trunc('Mon', ts), + date_trunc('week', ts) +FROM test_trunc_ts +ORDER BY ts + +-- NULL format is Incompatible on the native path. Without allowIncompatible the +-- codegen dispatcher runs Spark's TruncTimestamp and returns NULL. +query +SELECT ts, date_trunc(NULL, ts) FROM test_trunc_ts ORDER BY ts + +query +SELECT date_trunc('YEAR', NULL), date_trunc(NULL, NULL) + +-- Literal arguments exercise the scalar timestamp input branch. +query +SELECT + date_trunc('YEAR', TIMESTAMP '2024-05-17 12:34:56.123456'), + date_trunc('QUARTER', TIMESTAMP '2024-05-17 12:34:56.123456'), + date_trunc('MONTH', TIMESTAMP '2024-05-17 12:34:56.123456'), + date_trunc('WEEK', TIMESTAMP '2024-05-17 12:34:56.123456'), + date_trunc('DAY', TIMESTAMP '2024-05-17 12:34:56.123456'), + date_trunc('HOUR', TIMESTAMP '2024-05-17 12:34:56.123456'), + date_trunc('MINUTE', TIMESTAMP '2024-05-17 12:34:56.123456'), + date_trunc('SECOND', TIMESTAMP '2024-05-17 12:34:56.123456'), + date_trunc('MILLISECOND', TIMESTAMP '2024-05-17 12:34:56.123456'), + date_trunc('MICROSECOND', TIMESTAMP '2024-05-17 12:34:56.123456') + +-- Literal arguments exercise the scalar timestamp input branch. +query +SELECT + date_trunc('YEAR', TIMESTAMP '3333-05-17 12:34:56.123456'), + date_trunc('QUARTER', TIMESTAMP '3333-05-17 12:34:56.123456'), + date_trunc('MONTH', TIMESTAMP '3333-05-17 12:34:56.123456'), + date_trunc('WEEK', TIMESTAMP '3333-05-17 12:34:56.123456'), + date_trunc('DAY', TIMESTAMP '3333-05-17 12:34:56.123456'), + date_trunc('HOUR', TIMESTAMP '3333-05-17 12:34:56.123456'), + date_trunc('MINUTE', TIMESTAMP '3333-05-17 12:34:56.123456'), + date_trunc('SECOND', TIMESTAMP '3333-05-17 12:34:56.123456'), + date_trunc('MILLISECOND', TIMESTAMP '3333-05-17 12:34:56.123456'), + date_trunc('MICROSECOND', TIMESTAMP '3333-05-17 12:34:56.123456') + +-- Long.MaxValue is used as an end-of-time sentinel and exceeds chrono's range. +-- Use microsecond storage to preserve the extreme value without an INT96 conversion. +-- Direct date_trunc projections must execute natively and match Spark exactly. +-- Wrapping them in unix_micros would dispatch the whole subtree to the JVM. +statement +CREATE TABLE test_trunc_ts_extreme(ts timestamp) USING parquet + +statement +INSERT INTO test_trunc_ts_extreme VALUES + (timestamp_micros(9223372036854775807)), + (timestamp_micros(-9000000000000000000)), + (timestamp('2024-05-17 12:34:56.123456')), + (timestamp('3333-05-17 12:34:56.123456')), + (NULL) + +query expect_native(date_trunc) +SELECT + date_trunc('YEAR', ts), + date_trunc('QUARTER', ts), + date_trunc('MONTH', ts), + date_trunc('WEEK', ts) +FROM test_trunc_ts_extreme +ORDER BY ts + +-- A valid microsecond input can truncate below Long.MinValue. Spark raises an error. +statement +CREATE TABLE test_trunc_ts_overflow(ts timestamp) USING parquet + +statement +INSERT INTO test_trunc_ts_overflow VALUES (timestamp_micros(-9223372036854775808)) + +query expect_error(long overflow) +SELECT date_trunc('YEAR', ts) FROM test_trunc_ts_overflow + +query expect_error(long overflow) +SELECT date_trunc('QUARTER', ts) FROM test_trunc_ts_overflow + +query expect_error(long overflow) +SELECT date_trunc('MONTH', ts) FROM test_trunc_ts_overflow + +query expect_error(long overflow) +SELECT date_trunc('WEEK', ts) FROM test_trunc_ts_overflow + +query expect_error(long overflow) +SELECT date_trunc('MINUTE', ts) FROM test_trunc_ts_overflow + +query expect_error(long overflow) +SELECT date_trunc('HOUR', ts) FROM test_trunc_ts_overflow + +query expect_error(long overflow) +SELECT date_trunc('DAY', ts) FROM test_trunc_ts_overflow + +-- SECOND/MILLISECOND at the lower bound depend on the Spark version. See +-- trunc_timestamp_fine_overflow.sql and trunc_timestamp_fine_overflow_spark42.sql. diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_dst_ambiguous.sql b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_dst_ambiguous.sql index a1059736b9c..74d2c09ab4c 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_dst_ambiguous.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_dst_ambiguous.sql @@ -15,36 +15,62 @@ -- specific language governing permissions and limitations -- under the License. --- Regression test for ambiguous local times during DST fall-back. --- On 2024-11-03 at 2:00 AM America/Los_Angeles, clocks fall back to 1:00 AM, --- so 1:30 AM occurs twice (once in PDT, once in PST). Truncating 01:30 to HOUR --- gives 01:00, which is ambiguous. chrono's DateTime::with_minute(0) returns --- None for ambiguous results, causing a panic in as_micros_from_unix_epoch_utc. +-- Differential coverage for scalar date_trunc around DST overlaps, gaps, and historical offsets +-- containing seconds. Explicit UTC offsets include both occurrences of the repeated US fall-back +-- hour. The 2018 Sao Paulo values exercise its historic midnight spring-forward gap: truncating a +-- valid 01:30 local timestamp to DAY targets the nonexistent local midnight. Africa/Monrovia used +-- UTC-00:44:30 until 1972, so its local minute boundaries do not align with UTC minute boundaries. +-- Asia/Aden covers minute truncation across a historical offset transition. Havana repeated +-- midnight on 2020-11-01, which makes MONTH select the earlier UTC occurrence. Toronto +-- skipped 1919-03-30 23:30 through 1919-03-31 00:30, so WEEK must select the gap end. +-- Asuncion skipped midnight on 2023-10-01, a MONTH and QUARTER boundary. -- Config: spark.comet.expression.TruncTimestamp.allowIncompatible=true --- Config: spark.sql.session.timeZone=America/Los_Angeles +-- Config: spark.sql.parquet.int96RebaseModeInWrite=CORRECTED +-- ConfigMatrix: spark.sql.session.timeZone=America/Los_Angeles,America/New_York,America/Sao_Paulo,Africa/Monrovia,Asia/Aden,America/Havana,America/Toronto,America/Asuncion statement CREATE TABLE test_trunc_ambiguous(ts timestamp) USING parquet statement INSERT INTO test_trunc_ambiguous VALUES - (timestamp('2024-11-03 01:30:00')) + (timestamp('1919-04-02T16:00:00Z')), + (timestamp('1919-03-31T04:30:00Z')), + (timestamp('2023-10-15T12:00:00Z')), + (TIMESTAMP '1960-06-15 10:30:45'), + (timestamp('1947-03-13T20:53:30.123Z')), + (timestamp('1972-01-07T00:44:45Z')), + (timestamp('2018-11-04T01:30:15.123456Z')), + (timestamp('2018-11-04T02:30:15.123456Z')), + (timestamp('2018-11-04T03:30:15.123456Z')), + (timestamp('2018-11-04T04:30:15.123456Z')), + (timestamp('2024-03-10T06:30:15.123456Z')), + (timestamp('2024-03-10T07:30:15.123456Z')), + (timestamp('2024-03-10T08:30:15.123456Z')), + (timestamp('2024-03-10T09:30:15.123456Z')), + (timestamp('2024-03-10T10:30:15.123456Z')), + (timestamp('2024-03-10T11:30:15.123456Z')), + (timestamp('2024-11-03T05:30:15.123456Z')), + (timestamp('2024-11-03T06:30:15.123456Z')), + (timestamp('2024-11-03T07:30:15.123456Z')), + (timestamp('2024-11-03T08:30:15.123456Z')), + (timestamp('2024-11-03T09:30:15.123456Z')), + (timestamp('2024-11-03T10:30:15.123456Z')), + (timestamp('2020-11-15T12:00:00Z')), + (NULL) -query ignore(native panic: chrono returns None for ambiguous local time during DST fall-back) -SELECT ts, date_trunc('DAY', ts) FROM test_trunc_ambiguous ORDER BY ts - -query ignore(native panic: chrono returns None for ambiguous local time during DST fall-back) -SELECT ts, date_trunc('HOUR', ts) FROM test_trunc_ambiguous ORDER BY ts - -query ignore(native panic: chrono returns None for ambiguous local time during DST fall-back) -SELECT ts, date_trunc('WEEK', ts) FROM test_trunc_ambiguous ORDER BY ts - -query ignore(native panic: chrono returns None for ambiguous local time during DST fall-back) -SELECT ts, date_trunc('MONTH', ts) FROM test_trunc_ambiguous ORDER BY ts - -query ignore(native panic: chrono returns None for ambiguous local time during DST fall-back) -SELECT ts, date_trunc('QUARTER', ts) FROM test_trunc_ambiguous ORDER BY ts - -query ignore(native panic: chrono returns None for ambiguous local time during DST fall-back) -SELECT ts, date_trunc('YEAR', ts) FROM test_trunc_ambiguous ORDER BY ts +query +SELECT + ts, + date_trunc('YEAR', ts), + date_trunc('QUARTER', ts), + date_trunc('MONTH', ts), + date_trunc('WEEK', ts), + date_trunc('DAY', ts), + date_trunc('HOUR', ts), + date_trunc('MINUTE', ts), + date_trunc('SECOND', ts), + date_trunc('MILLISECOND', ts), + date_trunc('MICROSECOND', ts) +FROM test_trunc_ambiguous +ORDER BY ts diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_dst_format_column.sql b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_dst_format_column.sql new file mode 100644 index 00000000000..e9533ae3cc4 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_dst_format_column.sql @@ -0,0 +1,74 @@ +-- 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. + + +-- date_trunc with the format in a column, around DST overlaps, gaps and historical offsets. Each +-- row is truncated with the rules a literal format applies (see trunc_timestamp_dst_ambiguous.sql +-- for what each zone covers), and a NULL format gives NULL. A non-literal format is Incompatible, +-- because an invalid format throws instead of returning NULL. + +-- Config: spark.comet.expression.TruncTimestamp.allowIncompatible=true +-- Config: spark.sql.parquet.int96RebaseModeInWrite=CORRECTED +-- ConfigMatrix: spark.sql.session.timeZone=UTC,America/Los_Angeles,America/New_York,America/Sao_Paulo,Africa/Monrovia,Asia/Aden,America/Havana,America/Toronto,America/Asuncion + +statement +CREATE TABLE test_trunc_fmt_ts(ts timestamp) USING parquet + +statement +INSERT INTO test_trunc_fmt_ts VALUES + (timestamp('1919-04-02T16:00:00Z')), + (timestamp('1919-03-31T04:30:00Z')), + (timestamp('1919-03-31T04:45:00Z')), + (timestamp('2023-10-15T12:00:00Z')), + (TIMESTAMP '1960-06-15 10:30:45'), + (timestamp('1947-03-13T20:53:30.123Z')), + (timestamp('1972-01-07T00:44:45Z')), + (timestamp('2018-11-04T02:30:15.123456Z')), + (timestamp('2018-11-04T03:30:15.123456Z')), + (timestamp('2024-03-10T09:30:15.123456Z')), + (timestamp('2024-03-10T10:30:15.123456Z')), + (timestamp('2024-11-03T08:30:15.123456Z')), + (timestamp('2024-11-03T09:30:15.123456Z')), + (timestamp('2020-11-01T04:30:00Z')), + (timestamp('2020-11-01T05:30:00Z')), + (timestamp('2020-11-15T12:00:00Z')), + (NULL) + +statement +CREATE TABLE test_trunc_fmt_unit(fmt string) USING parquet + +statement +INSERT INTO test_trunc_fmt_unit VALUES + ('YEAR'), ('yyyy'), ('YY'), ('QUARTER'), ('MONTH'), ('mon'), ('MM'), ('WEEK'), ('DAY'), ('dd'), + ('HOUR'), ('MINUTE'), ('SECOND'), ('MILLISECOND'), ('MICROSECOND'), (NULL) + +statement +CREATE TABLE test_trunc_fmt(ts timestamp, fmt string) USING parquet + +statement +INSERT INTO test_trunc_fmt SELECT ts, fmt FROM test_trunc_fmt_ts CROSS JOIN test_trunc_fmt_unit + +query expect_native(date_trunc) +SELECT ts, fmt, date_trunc(fmt, ts) FROM test_trunc_fmt ORDER BY ts, fmt + +-- A NULL literal format gives NULL. +query expect_native(date_trunc) +SELECT ts, date_trunc(NULL, ts) FROM test_trunc_fmt_ts ORDER BY ts + +-- A literal timestamp with the format in a column. 2024-11-03 01:30 is in the US fall-back overlap. +query expect_native(date_trunc) +SELECT fmt, date_trunc(fmt, TIMESTAMP '2024-11-03 01:30:00') FROM test_trunc_fmt_unit ORDER BY fmt diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_fine_overflow.sql b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_fine_overflow.sql new file mode 100644 index 00000000000..4ef39e12a95 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_fine_overflow.sql @@ -0,0 +1,56 @@ +-- 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. + +-- Before 4.2, Spark wraps the Long subtraction when SECOND or MILLISECOND truncation falls below +-- Long.MinValue. Spark 4.2 and later (SPARK-56663) raise instead, which +-- trunc_timestamp_fine_overflow_spark42.sql covers. +-- MaxSparkVersion: 4.1 + +-- Use microsecond storage to preserve the extreme value without an INT96 conversion. +-- Config: spark.sql.session.timeZone=UTC +-- A format column is Incompatible, because an invalid format throws instead of returning NULL. +-- Config: spark.comet.expression.TruncTimestamp.allowIncompatible=true +-- Config: spark.sql.parquet.int96RebaseModeInWrite=CORRECTED +-- Config: spark.sql.parquet.datetimeRebaseModeInWrite=CORRECTED +-- Config: spark.sql.parquet.outputTimestampType=TIMESTAMP_MICROS +-- ConfigMatrix: parquet.enable.dictionary=false,true + +statement +CREATE TABLE test_trunc_ts_fine_overflow(ts timestamp) USING parquet + +-- Long.MinValue, and an ordinary value. +statement +INSERT INTO test_trunc_ts_fine_overflow VALUES + (timestamp_micros(-9223372036854775808)), + (timestamp('2024-05-17 12:34:56.123456')) + +query expect_native(date_trunc) +SELECT date_trunc('SECOND', ts), date_trunc('MILLISECOND', ts) +FROM test_trunc_ts_fine_overflow + +statement +CREATE TABLE test_trunc_ts_fine_overflow_fmt(ts timestamp, fmt string) USING parquet + +statement +INSERT INTO test_trunc_ts_fine_overflow_fmt VALUES + (timestamp_micros(-9223372036854775808), 'SECOND'), + (timestamp_micros(-9223372036854775808), 'MILLISECOND'), + (timestamp('2024-05-17 12:34:56.123456'), 'SECOND') + +-- The same truncation with the format in a column. +query expect_native(date_trunc) +SELECT fmt, date_trunc(fmt, ts) FROM test_trunc_ts_fine_overflow_fmt ORDER BY fmt, ts diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_fine_overflow_spark42.sql b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_fine_overflow_spark42.sql new file mode 100644 index 00000000000..d6fadeeb288 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_fine_overflow_spark42.sql @@ -0,0 +1,71 @@ +-- 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. + +-- Spark 4.2 and later (SPARK-56663) raise when SECOND or MILLISECOND truncation falls below +-- Long.MinValue, like every other unit. trunc_timestamp_fine_overflow.sql covers 4.1 and +-- earlier, where Spark wraps the Long subtraction instead. +-- MinSparkVersion: 4.2 + +-- Use microsecond storage to preserve the extreme value without an INT96 conversion. +-- Config: spark.sql.session.timeZone=UTC +-- A format column is Incompatible, because an invalid format throws instead of returning NULL. +-- Config: spark.comet.expression.TruncTimestamp.allowIncompatible=true +-- Config: spark.sql.parquet.int96RebaseModeInWrite=CORRECTED +-- Config: spark.sql.parquet.datetimeRebaseModeInWrite=CORRECTED +-- Config: spark.sql.parquet.outputTimestampType=TIMESTAMP_MICROS +-- ConfigMatrix: parquet.enable.dictionary=false,true + +statement +CREATE TABLE test_trunc_ts_fine_overflow(ts timestamp) USING parquet + +-- Long.MinValue, and an ordinary value. +statement +INSERT INTO test_trunc_ts_fine_overflow VALUES + (timestamp_micros(-9223372036854775808)), + (timestamp('2024-05-17 12:34:56.123456')) + +-- Shows date_trunc runs natively here, so the errors below come from Comet, not a fallback. +query expect_native(date_trunc) +SELECT date_trunc('SECOND', ts), date_trunc('MILLISECOND', ts) +FROM test_trunc_ts_fine_overflow +WHERE ts > timestamp_micros(-9223372036854775808) + +query expect_error(long overflow) +SELECT date_trunc('SECOND', ts) FROM test_trunc_ts_fine_overflow + +query expect_error(long overflow) +SELECT date_trunc('MILLISECOND', ts) FROM test_trunc_ts_fine_overflow + +statement +CREATE TABLE test_trunc_ts_fine_overflow_fmt(ts timestamp, fmt string) USING parquet + +statement +INSERT INTO test_trunc_ts_fine_overflow_fmt VALUES + (timestamp_micros(-9223372036854775808), 'SECOND'), + (timestamp_micros(-9223372036854775808), 'MILLISECOND'), + (timestamp('2024-05-17 12:34:56.123456'), 'SECOND') + +-- With the format in a column. The first query shows it runs natively. +query expect_native(date_trunc) +SELECT date_trunc(fmt, ts) FROM test_trunc_ts_fine_overflow_fmt +WHERE ts > timestamp_micros(-9223372036854775808) + +query expect_error(long overflow) +SELECT date_trunc(fmt, ts) FROM test_trunc_ts_fine_overflow_fmt + +query expect_error(long overflow) +SELECT date_trunc(fmt, ts) FROM test_trunc_ts_fine_overflow_fmt WHERE fmt = 'MILLISECOND' diff --git a/spark/src/test/resources/sql-tests/expressions/map/explode.sql b/spark/src/test/resources/sql-tests/expressions/map/explode.sql new file mode 100644 index 00000000000..6a47e40aa96 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/map/explode.sql @@ -0,0 +1,142 @@ +-- 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. + +-- Config: spark.comet.exec.explode.enabled=true + +statement +CREATE TABLE explode_maps(id int, m map) USING parquet + +statement +INSERT INTO explode_maps VALUES + (1, map('b', 20, 'a', 10)), (2, map()), (3, NULL), (4, map('null', NULL)) + +query +SELECT id, explode(m) FROM explode_maps + +query +SELECT id, m, explode_outer(m) FROM explode_maps + +query +SELECT id, posexplode(m) FROM explode_maps + +query +SELECT id, m, posexplode_outer(m) FROM explode_maps + +query +SELECT posexplode_outer(m) FROM explode_maps + +query +SELECT id, k, v FROM explode_maps LATERAL VIEW OUTER explode(m) e AS k, v +WHERE k IS NULL OR v IS NULL + +query +SELECT id, p, k, v FROM explode_maps LATERAL VIEW OUTER posexplode(m) e AS p, k, v +WHERE p = 0 OR p IS NULL + +-- Native map producers feed the generator directly, without a Parquet map column. +statement +CREATE TABLE explode_map_inputs( + id int, k array, v array, + entries array>, text string) USING parquet + +statement +INSERT INTO explode_map_inputs VALUES + (1, array(2, 1), array('b', 'a'), array(struct(2, 'b'), struct(1, 'a')), 'b:2,a:1'), + (2, array(), array(), array(), ''), + (3, NULL, NULL, NULL, NULL), + (4, array(3), array(NULL), array(struct(3, NULL)), 'missing') + +query expect_native(map_from_arrays) +SELECT id, explode(map_from_arrays(k, v)) FROM explode_map_inputs + +query expect_native(map_from_arrays) +SELECT id, explode_outer(map_from_arrays(k, v)) FROM explode_map_inputs + +query expect_native(map_from_arrays) +SELECT id, posexplode(map_from_arrays(k, v)) FROM explode_map_inputs + +query expect_native(map_from_arrays) +SELECT id, posexplode_outer(map_from_arrays(k, v)) FROM explode_map_inputs + +query expect_native(map_from_entries) +SELECT id, explode(map_from_entries(entries)) FROM explode_map_inputs + +query expect_native(map_from_entries) +SELECT id, explode_outer(map_from_entries(entries)) FROM explode_map_inputs + +query expect_native(map_from_entries) +SELECT id, posexplode(map_from_entries(entries)) FROM explode_map_inputs + +query expect_native(map_from_entries) +SELECT id, posexplode_outer(map_from_entries(entries)) FROM explode_map_inputs + +query expect_native(str_to_map) +SELECT id, explode(str_to_map(text)) FROM explode_map_inputs + +query expect_native(str_to_map) +SELECT id, explode_outer(str_to_map(text)) FROM explode_map_inputs + +query expect_native(str_to_map) +SELECT id, posexplode(str_to_map(text)) FROM explode_map_inputs + +query expect_native(str_to_map) +SELECT id, posexplode_outer(str_to_map(text)) FROM explode_map_inputs + +-- Computed and scalar maps, including non-nullable values and nested fields. +query +SELECT id, posexplode(map(id, named_struct('v', coalesce(id, 0)))) FROM explode_maps + +query +SELECT id, explode(map(1, 10, 2, 20)) FROM explode_maps + +query +SELECT id, posexplode_outer(cast(NULL AS map)) FROM explode_maps + +-- The untyped map() constructor still falls back before the cast. +query expect_fallback(MapType(NullType,NullType,false)) +SELECT id, explode_outer(cast(map() AS map)) FROM explode_maps + +statement +CREATE TABLE explode_nested_maps( + id int, m map, struct, m: map>>) USING parquet + +statement +INSERT INTO explode_nested_maps VALUES + (1, map(array(2, 1), named_struct('a', array(10, NULL), 'm', map('x', unhex('FF00'))), + array(3), NULL)), + (2, map()), (3, NULL) + +query +SELECT id, explode(m) FROM explode_nested_maps + +query +SELECT id, explode_outer(m) FROM explode_nested_maps + +query +SELECT id, posexplode(m) FROM explode_nested_maps + +query +SELECT id, posexplode_outer(m) FROM explode_nested_maps + +-- Project the output fields and compose map expansion with array expansion. +query +SELECT id, k, v.a, v.m FROM explode_nested_maps LATERAL VIEW OUTER explode(m) e AS k, v + +query +SELECT id, p, k, item FROM explode_nested_maps +LATERAL VIEW OUTER posexplode(m) e AS p, k, v +LATERAL VIEW OUTER explode(v.a) a AS item diff --git a/spark/src/test/resources/sql-tests/expressions/map/explode_float_keys.sql b/spark/src/test/resources/sql-tests/expressions/map/explode_float_keys.sql new file mode 100644 index 00000000000..58b04c51f98 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/map/explode_float_keys.sql @@ -0,0 +1,47 @@ +-- 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. + +-- Config: spark.comet.exec.explode.enabled=true +-- Config: spark.sql.legacy.disableMapKeyNormalization=true + +-- Preserve -0.0 when writing the fixture on Spark 4.0+. The generator must return +-- the stored keys unchanged. Native map construction's normalization is tracked +-- separately in https://github.com/apache/datafusion-comet/issues/6549. +statement +CREATE TABLE explode_float_maps(id int, m map) USING parquet + +statement +INSERT INTO explode_float_maps VALUES + (1, map(cast('-0.0' AS double), 10, cast('NaN' AS double), 20)), + (2, map(cast('NaN' AS double), NULL)), (3, map()), (4, NULL) + +-- Render keys as strings too, so a comparison cannot silently equate the two zero signs. +query +SELECT id, cast(key AS string), value +FROM explode_float_maps LATERAL VIEW explode(m) e AS key, value + +query +SELECT id, cast(key AS string), value +FROM explode_float_maps LATERAL VIEW OUTER explode(m) e AS key, value + +query +SELECT id, pos, cast(key AS string), value +FROM explode_float_maps LATERAL VIEW posexplode(m) e AS pos, key, value + +query +SELECT id, pos, cast(key AS string), value +FROM explode_float_maps LATERAL VIEW OUTER posexplode(m) e AS pos, key, value diff --git a/spark/src/test/resources/sql-tests/expressions/math/abs_null_literal.sql b/spark/src/test/resources/sql-tests/expressions/math/abs_null_literal.sql new file mode 100644 index 00000000000..1f8a03b0929 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/math/abs_null_literal.sql @@ -0,0 +1,29 @@ +-- 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. + +-- ExcludeRules: org.apache.spark.sql.catalyst.optimizer.NullPropagation + +statement +CREATE TABLE abs_null_literal(v INT) USING parquet + +statement +INSERT INTO abs_null_literal VALUES (1), (-2), (NULL) + +-- Excluding ConstantFolding alone is insufficient: NullPropagation would replace +-- abs(NULL) with a NULL literal. Pin abs itself so that such a rewrite fails this test. +query expect_native(abs) +SELECT abs(CAST(NULL AS INT)) FROM abs_null_literal diff --git a/spark/src/test/resources/sql-tests/expressions/math/arithmetic.sql b/spark/src/test/resources/sql-tests/expressions/math/arithmetic.sql index 969ddf18511..a2cb88d3433 100644 --- a/spark/src/test/resources/sql-tests/expressions/math/arithmetic.sql +++ b/spark/src/test/resources/sql-tests/expressions/math/arithmetic.sql @@ -79,3 +79,7 @@ INSERT INTO test_div_zero VALUES (1, 0, 0.0), (0, 0, 0.0) query SELECT a / b, a % b, d / 0.0 FROM test_div_zero + +-- `-0.0D` is a double literal with the sign bit set, and is a zero divisor too +query +SELECT d / -0.0D, d % -0.0D, d / double('-0.0'), d % double('-0.0') FROM test_div_zero diff --git a/spark/src/test/resources/sql-tests/expressions/math/nan_divisor.sql b/spark/src/test/resources/sql-tests/expressions/math/nan_divisor.sql new file mode 100644 index 00000000000..51135545d44 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/math/nan_divisor.sql @@ -0,0 +1,52 @@ +-- 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. + +-- Division and remainder pass a NaN divisor through, so `1.0D / (-d)` of a NaN is a NaN with the +-- sign bit set, the NaN that arithmetic produces on x86-64. Spark hashes every NaN alike, and +-- orders every NaN above every other value. ANSI mode divides through a different native path. +-- ConfigMatrix: spark.sql.ansi.enabled=false,true + +statement +CREATE TABLE nan_divisor(d double, f float) USING parquet + +statement +INSERT INTO nan_divisor VALUES (double('NaN'), float('NaN')), (1.5, float(1.5)), (NULL, NULL) + +query +SELECT hash(1.0D / (-d)), xxhash64(1.0D / (-d)), hash(1.0D % (-d)), xxhash64(1.0D % (-d)) +FROM nan_divisor + +query +SELECT hash(1.0F % (-f)), xxhash64(1.0F % (-f)) FROM nan_divisor + +query +SELECT least(1.0D / (-d), 0.0D), greatest(1.0D / (-d), 0.0D), least(1.0D % (-d), 0.0D), + greatest(1.0F % (-f), 0.0F) +FROM nan_divisor + +-- The union keeps the quotient in a projection below the aggregate. +query +SELECT max(q), min(q) +FROM (SELECT 1.0D / (-d) AS q FROM nan_divisor UNION ALL SELECT 0.0D AS q FROM nan_divisor) + +-- percentile_approx orders a NaN with the sign bit set below every other value +-- (https://github.com/apache/datafusion-comet/issues/6519). Until that is fixed, Comet keeps +-- normalizing the divisor, which makes the quotient a canonical NaN. +query +SELECT percentile_approx(q, 1.0D), percentile_approx(q, 0.0D) +FROM (SELECT /*+ COALESCE(1) */ q + FROM (SELECT 1.0D / (-d) AS q FROM nan_divisor UNION ALL SELECT 0.0D AS q FROM nan_divisor)) diff --git a/spark/src/test/resources/sql-tests/expressions/string/regexp_extract.sql b/spark/src/test/resources/sql-tests/expressions/string/regexp_extract.sql index adea5b30ec5..81f01248444 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/regexp_extract.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/regexp_extract.sql @@ -16,6 +16,7 @@ -- under the License. -- Test regexp_extract via JVM regex engine +-- Raw literals keep Spark from dropping regex backslashes such as '\d'. statement CREATE TABLE test_regexp_extract(s string) USING parquet @@ -24,33 +25,33 @@ statement INSERT INTO test_regexp_extract VALUES ('abc123def'), ('no match'), (NULL), ('xyz789'), ('hello world'), ('aa') -- group 0: entire match -query -SELECT regexp_extract(s, '\d+', 0) FROM test_regexp_extract +query expect_dispatch(regexp_extract) +SELECT regexp_extract(s, r'\d+', 0) FROM test_regexp_extract -- group 1: first capturing group -query -SELECT regexp_extract(s, '([a-z]+)(\d+)', 1) FROM test_regexp_extract +query expect_dispatch(regexp_extract) +SELECT regexp_extract(s, r'([a-z]+)(\d+)', 1) FROM test_regexp_extract -- group 2: second capturing group -query -SELECT regexp_extract(s, '([a-z]+)(\d+)', 2) FROM test_regexp_extract +query expect_dispatch(regexp_extract) +SELECT regexp_extract(s, r'([a-z]+)(\d+)', 2) FROM test_regexp_extract -- no match returns empty string -query +query expect_dispatch(regexp_extract) SELECT regexp_extract(s, 'NOMATCH', 0) FROM test_regexp_extract -- backreference pattern (Java-only) -query -SELECT regexp_extract(s, '(\w)\1', 0) FROM test_regexp_extract +query expect_dispatch(regexp_extract) +SELECT regexp_extract(s, r'(\w)\1', 0) FROM test_regexp_extract -- lookahead (Java-only) -query -SELECT regexp_extract(s, 'abc(?=\d)', 0) FROM test_regexp_extract +query expect_dispatch(regexp_extract) +SELECT regexp_extract(s, r'abc(?=\d)', 0) FROM test_regexp_extract -- embedded flags (Java-only) -query +query expect_dispatch(regexp_extract) SELECT regexp_extract(s, '(?i)HELLO', 0) FROM test_regexp_extract -- literal arguments -query -SELECT regexp_extract('abc123', '(\d+)', 1), regexp_extract('no digits', '(\d+)', 1), regexp_extract(NULL, '(\d+)', 1) +query expect_dispatch(regexp_extract) +SELECT regexp_extract('abc123', r'(\d+)', 1), regexp_extract('no digits', r'(\d+)', 1), regexp_extract(NULL, r'(\d+)', 1) diff --git a/spark/src/test/resources/sql-tests/expressions/string/regexp_extract_all.sql b/spark/src/test/resources/sql-tests/expressions/string/regexp_extract_all.sql index 880c7515a28..af3ada87c9c 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/regexp_extract_all.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/regexp_extract_all.sql @@ -16,6 +16,7 @@ -- under the License. -- Test regexp_extract_all via JVM regex engine +-- Raw literals keep Spark from dropping regex backslashes such as '\d'. statement CREATE TABLE test_regexp_extract_all(s string) USING parquet @@ -24,29 +25,29 @@ statement INSERT INTO test_regexp_extract_all VALUES ('abc123def456'), ('no match'), (NULL), ('100-200-300'), ('hello world') -- group 0: all entire matches -query -SELECT regexp_extract_all(s, '\d+', 0) FROM test_regexp_extract_all +query expect_dispatch(regexp_extract_all) +SELECT regexp_extract_all(s, r'\d+', 0) FROM test_regexp_extract_all -- group 1: first capturing group from each match -query -SELECT regexp_extract_all(s, '([a-z]+)(\d+)', 1) FROM test_regexp_extract_all +query expect_dispatch(regexp_extract_all) +SELECT regexp_extract_all(s, r'([a-z]+)(\d+)', 1) FROM test_regexp_extract_all -- group 2: second capturing group from each match -query -SELECT regexp_extract_all(s, '([a-z]+)(\d+)', 2) FROM test_regexp_extract_all +query expect_dispatch(regexp_extract_all) +SELECT regexp_extract_all(s, r'([a-z]+)(\d+)', 2) FROM test_regexp_extract_all -- no match returns empty array -query +query expect_dispatch(regexp_extract_all) SELECT regexp_extract_all(s, 'NOMATCH', 0) FROM test_regexp_extract_all -- backreference pattern (Java-only) -query -SELECT regexp_extract_all(s, '(\d)\1', 0) FROM test_regexp_extract_all +query expect_dispatch(regexp_extract_all) +SELECT regexp_extract_all(s, r'(\d)\1', 0) FROM test_regexp_extract_all -- embedded flags (Java-only) -query +query expect_dispatch(regexp_extract_all) SELECT regexp_extract_all(s, '(?i)[A-Z]+', 0) FROM test_regexp_extract_all -- literal arguments -query -SELECT regexp_extract_all('abc123def456', '(\d+)', 1), regexp_extract_all('no digits', '(\d+)', 1), regexp_extract_all(NULL, '(\d+)', 1) +query expect_dispatch(regexp_extract_all) +SELECT regexp_extract_all('abc123def456', r'(\d+)', 1), regexp_extract_all('no digits', r'(\d+)', 1), regexp_extract_all(NULL, r'(\d+)', 1) diff --git a/spark/src/test/resources/sql-tests/expressions/string/regexp_instr.sql b/spark/src/test/resources/sql-tests/expressions/string/regexp_instr.sql index d4754ef5c22..71cf1c9dab1 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/regexp_instr.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/regexp_instr.sql @@ -16,6 +16,7 @@ -- under the License. -- Test regexp_instr via JVM regex engine +-- Raw literals keep Spark from dropping regex backslashes such as '\d'. statement CREATE TABLE test_regexp_instr(s string) USING parquet @@ -24,25 +25,25 @@ statement INSERT INTO test_regexp_instr VALUES ('abc123def'), ('no match'), (NULL), ('123xyz'), ('hello world'), ('aa') -- basic: position of first digit sequence -query -SELECT regexp_instr(s, '\d+', 0) FROM test_regexp_instr +query expect_dispatch(regexp_instr) +SELECT regexp_instr(s, r'\d+', 0) FROM test_regexp_instr -- group 1 (still returns position of entire match per Spark semantics) -query -SELECT regexp_instr(s, '([a-z]+)(\d+)', 1) FROM test_regexp_instr +query expect_dispatch(regexp_instr) +SELECT regexp_instr(s, r'([a-z]+)(\d+)', 1) FROM test_regexp_instr -- no match returns 0 -query +query expect_dispatch(regexp_instr) SELECT regexp_instr(s, 'NOMATCH', 0) FROM test_regexp_instr -- backreference pattern (Java-only) -query -SELECT regexp_instr(s, '(\w)\1', 0) FROM test_regexp_instr +query expect_dispatch(regexp_instr) +SELECT regexp_instr(s, r'(\w)\1', 0) FROM test_regexp_instr -- embedded flags (Java-only) -query +query expect_dispatch(regexp_instr) SELECT regexp_instr(s, '(?i)HELLO', 0) FROM test_regexp_instr -- literal arguments -query -SELECT regexp_instr('abc123', '\d+', 0), regexp_instr('no digits', '\d+', 0), regexp_instr(NULL, '\d+', 0) +query expect_dispatch(regexp_instr) +SELECT regexp_instr('abc123', r'\d+', 0), regexp_instr('no digits', r'\d+', 0), regexp_instr(NULL, r'\d+', 0) diff --git a/spark/src/test/resources/sql-tests/expressions/string/regexp_replace_java.sql b/spark/src/test/resources/sql-tests/expressions/string/regexp_replace_java.sql index c0e4f955606..211740e10a2 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/regexp_replace_java.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/regexp_replace_java.sql @@ -16,6 +16,7 @@ -- under the License. -- Test regexp_replace via JVM regex engine +-- Raw literals keep Spark from dropping regex backslashes such as '\d'. statement CREATE TABLE test_regexp_replace_java(s string) USING parquet @@ -23,28 +24,28 @@ CREATE TABLE test_regexp_replace_java(s string) USING parquet statement INSERT INTO test_regexp_replace_java VALUES ('100-200'), ('abc'), (''), (NULL), ('phone 123-456-7890'), ('aabbcc') -query -SELECT regexp_replace(s, '\d+', 'X') FROM test_regexp_replace_java +query expect_dispatch(regexp_replace) +SELECT regexp_replace(s, r'\d+', 'X') FROM test_regexp_replace_java -query -SELECT regexp_replace(s, '\d+', 'X', 1) FROM test_regexp_replace_java +query expect_dispatch(regexp_replace) +SELECT regexp_replace(s, r'\d+', 'X', 1) FROM test_regexp_replace_java -- backreference in replacement -query -SELECT regexp_replace(s, '(\d+)-(\d+)', '$2-$1') FROM test_regexp_replace_java +query expect_dispatch(regexp_replace) +SELECT regexp_replace(s, r'(\d+)-(\d+)', '$2-$1') FROM test_regexp_replace_java -- backreference in pattern (Java-only) -query -SELECT regexp_replace(s, '(\w)\1', 'Z') FROM test_regexp_replace_java +query expect_dispatch(regexp_replace) +SELECT regexp_replace(s, r'(\w)\1', 'Z') FROM test_regexp_replace_java -- lookahead (Java-only) -query -SELECT regexp_replace(s, '\d+(?=-)', 'X') FROM test_regexp_replace_java +query expect_dispatch(regexp_replace) +SELECT regexp_replace(s, r'\d+(?=-)', 'X') FROM test_regexp_replace_java -- embedded flags (Java-only) -query +query expect_dispatch(regexp_replace) SELECT regexp_replace(s, '(?i)ABC', 'X') FROM test_regexp_replace_java -- literal arguments -query -SELECT regexp_replace('100-200', '(\d+)', 'X'), regexp_replace('abc', '(\d+)', 'X'), regexp_replace(NULL, '(\d+)', 'X') +query expect_dispatch(regexp_replace) +SELECT regexp_replace('100-200', r'(\d+)', 'X'), regexp_replace('abc', r'(\d+)', 'X'), regexp_replace(NULL, r'(\d+)', 'X') diff --git a/spark/src/test/resources/sql-tests/expressions/string/regexp_replace_rust.sql b/spark/src/test/resources/sql-tests/expressions/string/regexp_replace_rust.sql index 6aabf1de5ef..2f5c4c94294 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/regexp_replace_rust.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/regexp_replace_rust.sql @@ -17,6 +17,7 @@ -- Test regexp_replace() with Rust regexp engine -- Opt in to the native (rust) regexp engine. Without this, it runs through the codegen dispatcher. +-- Raw literals keep Spark from dropping regex backslashes such as '\d'. -- Config: spark.comet.expression.RegExpReplace.allowIncompatible=true statement @@ -25,12 +26,12 @@ CREATE TABLE test_regexp_replace_enabled(s string) USING parquet statement INSERT INTO test_regexp_replace_enabled VALUES ('100-200'), ('abc'), (''), (NULL), ('phone 123-456-7890') -query -SELECT regexp_replace(s, '(\d+)', 'X') FROM test_regexp_replace_enabled +query expect_native(regexp_replace) +SELECT regexp_replace(s, r'(\d+)', 'X') FROM test_regexp_replace_enabled -query -SELECT regexp_replace(s, '(\d+)', 'X', 1) FROM test_regexp_replace_enabled +query expect_native(regexp_replace) +SELECT regexp_replace(s, r'(\d+)', 'X', 1) FROM test_regexp_replace_enabled -- literal + literal + literal -query -SELECT regexp_replace('100-200', '(\d+)', 'X'), regexp_replace('abc', '(\d+)', 'X'), regexp_replace(NULL, '(\d+)', 'X') +query expect_native(regexp_replace) +SELECT regexp_replace('100-200', r'(\d+)', 'X'), regexp_replace('abc', r'(\d+)', 'X'), regexp_replace(NULL, r'(\d+)', 'X') diff --git a/spark/src/test/resources/sql-tests/expressions/string/rlike_auto_native.sql b/spark/src/test/resources/sql-tests/expressions/string/rlike_auto_native.sql index e93b18491e9..4fe2792b8eb 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/rlike_auto_native.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/rlike_auto_native.sql @@ -16,8 +16,8 @@ -- under the License. -- Default settings: in-subset literal patterns run on the native Rust path. --- Routing is asserted in CometRegExpJvmSuite / CometRegexParitySuite; this file --- only checks result equality with Spark. +-- Each query checks both native routing and result equality with Spark. +-- Raw literals keep Spark from dropping regex backslashes such as '\d'. statement CREATE TABLE test_rlike_auto(s string) USING parquet @@ -25,20 +25,20 @@ CREATE TABLE test_rlike_auto(s string) USING parquet statement INSERT INTO test_rlike_auto VALUES ('hello'), ('12345'), (''), (NULL), ('Hello World'), ('abc123'), ('aa'), ('ab'), ('foo'), ('bar'), ('a+b') -query +query expect_native(rlike) SELECT s RLIKE 'abc[0-9]+' FROM test_rlike_auto -query +query expect_native(rlike) SELECT s RLIKE '[a-zA-Z_][a-zA-Z0-9_]*' FROM test_rlike_auto -query +query expect_native(rlike) SELECT s RLIKE '(foo|bar){1,3}' FROM test_rlike_auto -query -SELECT s RLIKE 'a\+b' FROM test_rlike_auto +query expect_native(rlike) +SELECT s RLIKE r'a\+b' FROM test_rlike_auto -query +query expect_native(rlike) SELECT s RLIKE '' FROM test_rlike_auto -query +query expect_native(rlike) SELECT 'hello' RLIKE '[a-z]+', '12345' RLIKE '[0-9]+', '' RLIKE '', NULL RLIKE 'a' diff --git a/spark/src/test/resources/sql-tests/expressions/string/rlike_java.sql b/spark/src/test/resources/sql-tests/expressions/string/rlike_java.sql index bf86973b232..7ea90ba90dc 100644 --- a/spark/src/test/resources/sql-tests/expressions/string/rlike_java.sql +++ b/spark/src/test/resources/sql-tests/expressions/string/rlike_java.sql @@ -15,7 +15,8 @@ -- specific language governing permissions and limitations -- under the License. --- Test RLIKE via JVM regex engine +-- Test RLIKE dispatcher patterns and the native empty-pattern case +-- Raw literals keep Spark from dropping regex backslashes such as '\d'. statement CREATE TABLE test_rlike_java(s string) USING parquet @@ -23,27 +24,27 @@ CREATE TABLE test_rlike_java(s string) USING parquet statement INSERT INTO test_rlike_java VALUES ('hello'), ('12345'), (''), (NULL), ('Hello World'), ('abc123'), ('aa'), ('ab') -query -SELECT s RLIKE '^\d+$' FROM test_rlike_java +query expect_dispatch(rlike) +SELECT s RLIKE r'^\d+$' FROM test_rlike_java -query +query expect_dispatch(rlike) SELECT s RLIKE '^[a-z]+$' FROM test_rlike_java -query +query expect_native(rlike) SELECT s RLIKE '' FROM test_rlike_java -- backreference (Java-only) -query -SELECT s RLIKE '^(\w)\1$' FROM test_rlike_java +query expect_dispatch(rlike) +SELECT s RLIKE r'^(\w)\1$' FROM test_rlike_java -- lookahead (Java-only) -query -SELECT s RLIKE 'abc(?=\d)' FROM test_rlike_java +query expect_dispatch(rlike) +SELECT s RLIKE r'abc(?=\d)' FROM test_rlike_java -- embedded flags (Java-only) -query +query expect_dispatch(rlike) SELECT s RLIKE '(?i)hello' FROM test_rlike_java --- literal arguments +-- Literal arguments mix native (empty pattern) and dispatched RLIKE, so only check parity. query -SELECT 'hello' RLIKE '^[a-z]+$', '12345' RLIKE '^\d+$', '' RLIKE '', NULL RLIKE 'a' +SELECT 'hello' RLIKE '^[a-z]+$', '12345' RLIKE r'^\d+$', '' RLIKE '', NULL RLIKE 'a' diff --git a/spark/src/test/resources/sql-tests/operators/one_row_relation.sql b/spark/src/test/resources/sql-tests/operators/one_row_relation.sql new file mode 100644 index 00000000000..e28d80cdf89 --- /dev/null +++ b/spark/src/test/resources/sql-tests/operators/one_row_relation.sql @@ -0,0 +1,44 @@ +-- 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. + +-- Queries without a FROM clause run over a OneRowRelation, which Comet converts by default. + +-- A folded array literal is an UnsafeArrayData, not a GenericArrayData. +query +SELECT sequence(1, 5) + +query +SELECT sequence(1, 5), array(array(1, 2), array(3)), array('a', NULL, 'c') + +query +SELECT explode(sequence(1, 3)) + +-- GROUPING SETS over a single row. +statement +CREATE TEMPORARY VIEW one_row_user_view AS SELECT current_user + +query +SELECT count(*) FROM one_row_user_view GROUP BY current_user GROUPING SETS ((current_user)) + +query +SELECT current_user, grouping(current_user) FROM one_row_user_view GROUP BY ROLLUP(current_user) + +query +SELECT count(*) FROM (SELECT 1 AS a) GROUP BY GROUPING SETS ((a), ()) + +statement +DROP VIEW one_row_user_view diff --git a/spark/src/test/resources/sql-tests/windows/nested_float_order_keys.sql b/spark/src/test/resources/sql-tests/windows/nested_float_order_keys.sql index 711cf63bc47..92e7d29f058 100644 --- a/spark/src/test/resources/sql-tests/windows/nested_float_order_keys.sql +++ b/spark/src/test/resources/sql-tests/windows/nested_float_order_keys.sql @@ -24,8 +24,8 @@ -- total order sorts below -Infinity. Each ORDER BY ends with a unique tiebreaker, so peers come -- out in the tiebreaker's order. --- Strict floating-point mode declines these keys, because their types can hold a null element or --- field: see nested_float_order_keys_strict.sql. +-- Strict floating-point mode runs these keys natively too. nested_float_order_keys_strict.sql covers +-- that, along with the null orders and RANGE frames that fall back in every mode. statement CREATE TABLE nested_float_keys(id INT, g INT, d DOUBLE, f FLOAT, s BOOLEAN) USING parquet diff --git a/spark/src/test/resources/sql-tests/windows/nested_float_order_keys_strict.sql b/spark/src/test/resources/sql-tests/windows/nested_float_order_keys_strict.sql index 995470b2273..3d83601fa8a 100644 --- a/spark/src/test/resources/sql-tests/windows/nested_float_order_keys_strict.sql +++ b/spark/src/test/resources/sql-tests/windows/nested_float_order_keys_strict.sql @@ -16,11 +16,11 @@ -- under the License. -- Strict floating-point mode and sort or window keys that nest floats in arrays and structs. --- nested_float_order_keys.sql checks that Comet normalizes those floats to match Spark. Strict --- mode still declines a key whose type can hold a null element or field, because Spark orders --- such a null below every value whatever the key's null order, the native sort places it by the --- null order (#6476), and a RANGE window frame orders it above every value (#6477). A key whose --- type cannot hold a null stays native. +-- nested_float_order_keys.sql checks that Comet normalizes those floats to match Spark, so strict +-- mode runs these keys natively too. Spark orders a null element or field below every other +-- value, whatever the key's null order. Comet falls back, in every mode, for the two shapes where +-- that differs: `ASC NULLS LAST` or `DESC NULLS FIRST` on a key whose type can hold such a null +-- (#6476), and a RANGE window frame that has to find a row's peers over such a key (#6477). -- -- The test harness admits incompatible sort orders by default, so turn that off to check the -- shipped policy. @@ -42,10 +42,11 @@ INSERT INTO nested_float_strict VALUES (8, 1, NULL, NULL, false), (9, 2, double('-Infinity'), float('-Infinity'), false) --- Keys over the nullable columns can hold a null element or field, whether or not any row does, --- so they fall back, and Spark's answers come back. Row 8 holds one. Native, ORDER BY array(d) --- NULLS LAST puts row 8 last, where Spark puts it first, and the running sum spans the whole --- partition on row 8 and every row after it. +-- Keys over the nullable columns can hold a null element or field, whether or not any row does. +-- Row 8 holds one. Under a non-default null order, or in a RANGE frame that finds a row's peers, +-- they fall back, and Spark's answers come back. Native, ORDER BY array(d) NULLS LAST would put +-- row 8 last, where Spark puts it first, and the running sum would span the whole partition on +-- row 8 and every row after it. query expect_fallback(can hold a null element or field) SELECT id FROM nested_float_strict ORDER BY array(d) NULLS LAST, id @@ -56,25 +57,51 @@ query expect_fallback(can hold a null element or field) SELECT id FROM nested_float_strict ORDER BY named_struct('x', f) DESC NULLS FIRST, id query expect_fallback(can hold a null element or field) +SELECT id, RANK() OVER (ORDER BY array(IF(s, -d, d)) DESC NULLS FIRST, id) AS r +FROM nested_float_strict + +-- The default null orders place a null element or field where Spark does, so the same keys run +-- natively in a Sort, a TopK, ranking functions, a rank limit, a ROWS frame and a RANGE frame +-- bounded only by the partition. Row 8's null element sorts first ascending and last descending. +-- IF(s, ...) makes row 1 a -0.0 and row 5 a NaN with the sign bit set. +query +SELECT id FROM nested_float_strict ORDER BY array(IF(s, -d, d)), id + +query +SELECT id FROM nested_float_strict ORDER BY named_struct('x', IF(s, -f, f)) DESC, id + +query SELECT id FROM nested_float_strict ORDER BY array(IF(s, -d, d)), id LIMIT 4 -query expect_fallback(can hold a null element or field) +query SELECT id, RANK() OVER (PARTITION BY g ORDER BY named_struct('x', IF(s, -d, d))) AS r, DENSE_RANK() OVER (ORDER BY array(IF(s, -f, f)) DESC) AS dr FROM nested_float_strict -query expect_fallback(can hold a null element or field) +query SELECT id FROM ( SELECT id, RANK() OVER (ORDER BY array(IF(s, -d, d))) AS r FROM nested_float_strict ) WHERE r <= 3 +query +SELECT id, + SUM(id) OVER (ORDER BY array(IF(s, -d, d)), id + ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS running, + SUM(id) OVER (PARTITION BY g ORDER BY named_struct('x', IF(s, -f, f)) + RANGE BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) AS total +FROM nested_float_strict + -- coalesce with a literal is never null, so an array or struct of it cannot hold a null, and --- these keys stay native. Row 8 becomes a third zero. Each ORDER BY ends with a unique --- tiebreaker, so peers come out in its order, and the running sums need no null row to leave out. +-- these keys stay native under any null order and in RANGE frames. Row 8 becomes a third zero. +-- Each ORDER BY ends with a unique tiebreaker, so peers come out in its order, and the running +-- sums need no null row to leave out. query SELECT id FROM nested_float_strict ORDER BY array(coalesce(IF(s, -d, d), 0.0D)), id DESC +query +SELECT id FROM nested_float_strict ORDER BY array(coalesce(IF(s, -d, d), 0.0D)) NULLS LAST, id + query SELECT id FROM nested_float_strict ORDER BY named_struct('x', coalesce(IF(s, -f, f), 0.0F)) DESC, id diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark4_2/q77a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark4_2/q77a/extended.txt index c97afc95a3f..30de4b735f5 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark4_2/q77a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark4_2/q77a/extended.txt @@ -65,7 +65,7 @@ CometColumnarToRow : : : +- CometProject : : : +- CometFilter : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : +- Project + : : +- Project [COMET: Unsupported data type: StructType(StructField(returns,DecimalType(17,2),true),StructField(profit_loss,DecimalType(17,2),true))] : : : :- Subquery : : : : +- CometColumnarToRow : : : : +- CometProject @@ -81,7 +81,7 @@ CometColumnarToRow : : : : +- CometFilter : : : : +- CometNativeScan parquet spark_catalog.default.date_dim : : : +- ReusedSubquery - : : +- Scan OneRowRelation [COMET: Scan OneRowRelation is not supported] + : : +- Scan OneRowRelation : +- CometColumnarToRow : +- CometProject : +- CometBroadcastHashJoin @@ -184,7 +184,7 @@ CometColumnarToRow : : : +- CometProject : : : +- CometFilter : : : +- CometNativeScan parquet spark_catalog.default.date_dim - : : +- Project + : : +- Project [COMET: Unsupported data type: StructType(StructField(returns,DecimalType(17,2),true),StructField(profit_loss,DecimalType(17,2),true))] : : : :- Subquery : : : : +- CometColumnarToRow : : : : +- CometProject @@ -200,7 +200,7 @@ CometColumnarToRow : : : : +- CometFilter : : : : +- CometNativeScan parquet spark_catalog.default.date_dim : : : +- ReusedSubquery - : : +- Scan OneRowRelation [COMET: Scan OneRowRelation is not supported] + : : +- Scan OneRowRelation : +- CometColumnarToRow : +- CometProject : +- CometBroadcastHashJoin @@ -303,7 +303,7 @@ CometColumnarToRow : : +- CometProject : : +- CometFilter : : +- CometNativeScan parquet spark_catalog.default.date_dim - : +- Project + : +- Project [COMET: Unsupported data type: StructType(StructField(returns,DecimalType(17,2),true),StructField(profit_loss,DecimalType(17,2),true))] : : :- Subquery : : : +- CometColumnarToRow : : : +- CometProject @@ -319,7 +319,7 @@ CometColumnarToRow : : : +- CometFilter : : : +- CometNativeScan parquet spark_catalog.default.date_dim : : +- ReusedSubquery - : +- Scan OneRowRelation [COMET: Scan OneRowRelation is not supported] + : +- Scan OneRowRelation +- CometColumnarToRow +- CometProject +- CometBroadcastHashJoin diff --git a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala index 5cbaf538316..c464746c8c7 100644 --- a/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala @@ -1379,12 +1379,10 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp } } - // The tests below deliberately live in this suite, not a `sql-tests` fixture: constant folding is - // enabled by default here, so each `map(...)` collapses to a MapType Literal and the outer - // `array(...)` reaches `CometCreateArray` with folded-Literal children, which `CometLiteral` - // rebuilds as an equivalent `CreateMap` of primitive literals -- the folded-literal expansion path - // under review. `CometSqlFileTestSuite` force-disables `ConstantFolding`, so an equivalent SQL - // fixture would only exercise the constructor path (which `create_array.sql` already covers). + // These tests exercise folded complex literals, including CometLiteral's expansion to native + // constructors. SQL fixtures need `-- ConstantFolding: enabled` to reach that same path. + // folded_array_map_literals.sql covers the mixed map/NULL array; the remaining conversions are + // tracked in #6628. test("array of folded map literals with array values (multirow)") { withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { checkSparkAnswerAndOperator( @@ -1432,17 +1430,6 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp } } - // A NULL element sits next to a populated one inside a single folded - // `ArrayType(MapType(IntegerType, IntegerType, false), containsNull = true)` literal. The null - // slot serializes as a typed null literal carrying the declared map type, so the rebuilt sibling - // has to report that same type for `make_array` to accept the pair. - test("folded array mixing a map literal and a NULL element (multirow)") { - withParquetTable((0 until 3).map(i => (i, i.toLong)), "tbl") { - checkSparkAnswerAndOperator("SELECT array(map(1, 2), NULL) AS arr FROM tbl") - checkSparkAnswerAndOperator("SELECT array(NULL, map(1, 2)) AS arr FROM tbl") - } - } - // A map value that is itself a map or an array of maps is built by Spark's own generated code // inside `CometCreateMap`, so it keeps its declared nullability all the way to the consumer. test("folded map literals with nested map values (multirow)") { diff --git a/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala index 99a57b15757..db72532ee50 100644 --- a/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometBitwiseExpressionSuite.scala @@ -54,26 +54,6 @@ class CometBitwiseExpressionSuite extends CometTestBase with AdaptiveSparkPlanHe } } - test("bitwise shift with different left/right types") { - Seq(false, true).foreach { dictionary => - withSQLConf("parquet.enable.dictionary" -> dictionary.toString) { - val table = "test" - withTable(table) { - sql(s"create table $table(col1 long, col2 int) using parquet") - sql(s"insert into $table values(1111, 2)") - sql(s"insert into $table values(1111, 2)") - sql(s"insert into $table values(3333, 4)") - sql(s"insert into $table values(5555, 6)") - - checkSparkAnswerAndOperator( - s"SELECT shiftright(col1, 2), shiftright(col1, col2) FROM $table") - checkSparkAnswerAndOperator( - s"SELECT shiftleft(col1, 2), shiftleft(col1, col2) FROM $table") - } - } - } - } - test("bitwise_get - throws exceptions") { def checkSparkAndCometEqualThrows(query: String): Unit = { checkSparkAnswerMaybeThrows(sql(query)) match { diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala index 400ec6fcd2f..54b59dec9dd 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala @@ -1436,8 +1436,8 @@ class CometCodegenSourceSuite extends AnyFunSuite { } test("closure-serialized bytes diverge for failOnError / roundOff / timeZoneId variants") { - // The dispatcher caches kernels by the closure-serialized bytes of the bound expression - // (`CacheKey.bytesKey`). For expression classes that carry a runtime-dependent boolean + // The dispatcher caches kernels by a digest of the closure-serialized bytes of the bound + // expression (`CacheKey.digest`). For expression classes that carry a runtime-dependent boolean // (failOnError, roundOff) or string (timeZoneId), two plan instances differing only in that // field must serialize to distinct byte sequences so they receive distinct cache entries. // A collision would let a kernel compiled for one variant (e.g. ANSI throw site) silently @@ -1523,20 +1523,20 @@ class CometCodegenSourceSuite extends AnyFunSuite { } test("CacheKey discriminates on ArrowColumnSpec.nullable") { - // Structural regression: same expression bytes and same Arrow vector class with different + // Structural regression: same expression digest and same Arrow vector class with different // `nullable` must produce non-equal cache keys. The dispatcher today hardcodes `nullable=true` // for top-level specs, so the two variants don't both arise from runtime data, but the case // class equality contract still has to discriminate so that any future tiered cache or test // construction can rely on it. The non-nullable variant's generated source emits a literal // `false` from `isNullAt`, distinct codegen output that we never want to silently share with // the nullable variant. - val bytes = java.nio.ByteBuffer.wrap(Array[Byte](1, 2, 3)) + val digest = java.nio.ByteBuffer.wrap(Array[Byte](1, 2, 3)) val nullable = IndexedSeq[ArrowColumnSpec](ArrowColumnSpec(varCharVectorClass, nullable = true)) val nonNullable = IndexedSeq[ArrowColumnSpec](ArrowColumnSpec(varCharVectorClass, nullable = false)) - val k1 = CometScalaUDFCodegen.CacheKey(bytes, nullable) - val k2 = CometScalaUDFCodegen.CacheKey(bytes, nonNullable) + val k1 = CometScalaUDFCodegen.CacheKey(0L, digest, nullable) + val k2 = CometScalaUDFCodegen.CacheKey(0L, digest, nonNullable) assert( k1 != k2, "expected nullable=true and nullable=false specs to produce distinct cache keys") diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index c9371fbf6da..248e903fbf7 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -28,7 +28,7 @@ import org.apache.spark.{SparkConf, SparkEnv, TaskContext} import org.apache.spark.sql.{CometTestBase, Row} import org.apache.spark.sql.api.java.UDF1 import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Add, Alias, ApplyFunctionExpression, AttributeReference, BoundReference, Cast, CreateArray, CreateMap, CreateNamedStruct, Expression, GenericInternalRow, Hypot, Literal, MapConcat, ScalaUDF} +import org.apache.spark.sql.catalyst.expressions.{Add, Alias, ApplyFunctionExpression, AttributeReference, AttributeSeq, BindReferences, BoundReference, Cast, CreateArray, CreateMap, CreateNamedStruct, Expression, GenericInternalRow, Hypot, If, IsNull, LessThan, Literal, MapConcat, Or, Rand, ScalaUDF} import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial} import org.apache.spark.sql.catalyst.expressions.objects.{Invoke, StaticInvoke} import org.apache.spark.sql.catalyst.util.GenericArrayData @@ -40,11 +40,12 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.{ByteArray, UTF8String} -import org.apache.comet.CometSparkSessionExtensions.isSpark41Plus +import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, isSpark41Plus} import org.apache.comet.codegen.CometBatchKernelCodegen import org.apache.comet.codegen.CometBatchKernelCodegen.ArrowColumnSpec import org.apache.comet.serde.{CometInvokeTargets, CometScalaUDF, QueryPlanSerde} import org.apache.comet.serde.ExprOuterClass.Expr.ExprStructCase +import org.apache.comet.udf.CometUdfBridge import org.apache.comet.udf.codegen.CometScalaUDFCodegen import org.apache.comet.vector.CometVector @@ -623,6 +624,30 @@ class CometCodegenSuite .contains(Set("hypot", "cast", "checkoverflow", "add"))) } + test("the serde ships a digest of the serialized expression at arg 0 (#6705)") { + // The dispatcher trusts a cache hit on the digest without comparing the bytes, so the digest + // the serde ships must be the digest of the bytes it ships next to it. Calling the dispatcher + // directly cannot catch a mismatch, because those tests build their own digest. + val x = AttributeReference("x", DoubleType, nullable = false)() + def payload(e: Expression): (Array[Byte], Array[Byte]) = { + val proto = QueryPlanSerde.exprToProto(e, Seq(x)).get + assert(proto.hasJvmScalarUdf) + val args = proto.getJvmScalarUdf.getArgsList + ( + args.get(0).getLiteral.getBytesVal.toByteArray, + args.get(1).getLiteral.getBytesVal.toByteArray) + } + + val (hypotDigest, hypotBytes) = payload(Hypot(x, Literal(4.0d))) + val (otherDigest, otherBytes) = payload(Hypot(x, Literal(5.0d))) + assert(hypotDigest.sameElements(CometScalaUDFCodegen.digest(hypotBytes))) + assert(otherDigest.sameElements(CometScalaUDFCodegen.digest(otherBytes))) + + // Expressions that differ only in a literal must not share a kernel. + assert(!hypotBytes.sameElements(otherBytes)) + assert(!hypotDigest.sameElements(otherDigest)) + } + test("tags copied onto the shared TrueLiteral do not leak into unrelated plans") { // Catalyst copies a rewritten node's tags onto its replacement, so a tagged expression that an // earlier query rewrote into `Literal.TrueLiteral` brands that process-wide singleton for the @@ -865,7 +890,7 @@ class CometCodegenSuite "same UDF over nullable and non-nullable columns gets distinct kernels with independent state") { // Two columns, same type, different schema-declared nullability. Same UDF applied to each // alongside a per-projection MonotonicallyIncreasingID. Each projection has its own MII - // child (different bytesKey), so each kernel must have its own counter advancing 0..N-1. + // child (a different digest), so each kernel must have its own counter advancing 0..N-1. // If the dispatcher collapses them onto one kernel or shares state somehow, the counters // would interleave and the output would diverge from Spark. spark.udf.register("withId", (s: String, id: Long) => s"${s}_${id}") @@ -1003,6 +1028,150 @@ class CometCodegenSuite } } + test("a primitive ScalaUDF's null guard runs in the UDF's kernel (#6704)") { + // Spark's `HandleNullInputsForUDF` wraps each call below as + // `if (isnull(a) or ...) null else f(knownnotnull(a), ...)`. Run natively, the `if` is a + // `CASE` that splits each batch on the predicate and merges the two halves back, so the guard + // goes to the kernel with the call, and neither `if` nor `isnull` is native. + spark.udf.register("plusOne", (x: Long) => x + 1) + spark.udf.register("combine", (a: Long, b: Int, s: String) => s"$s:${a + b}") + spark.udf.register("addBoth", (x: Long, y: Long) => x + y) + spark.udf.register("addThree", (x: Long, y: Int, z: Long) => x + y + z) + withTable("t") { + sql("CREATE TABLE t (a BIGINT, b INT, s STRING) USING parquet") + sql( + "INSERT INTO t VALUES (1, 10, 'x'), (NULL, 20, 'y'), (3, NULL, 'z'), " + + "(NULL, NULL, NULL), (-5, 7, NULL)") + val guard = Seq("if", "isnull") + checkSparkAnswerAndImpl(sql("SELECT plusOne(a) FROM t"), dispatched = "plusone" +: guard) + checkSparkAnswerAndImpl( + sql("SELECT max(plusOne(a)), count(plusOne(a)) FROM t"), + dispatched = "plusone" +: guard) + // An argument that is an expression is checked as it is. + checkSparkAnswerAndImpl( + sql("SELECT plusOne(a * 2) FROM t"), + dispatched = "plusone" +: guard) + // One `isnull` per primitive parameter, joined by `or`. The `String` one is not checked. + checkSparkAnswerAndImpl( + sql("SELECT combine(a, b, s) FROM t"), + dispatched = Seq("combine", "or") ++ guard) + // The optimizer drops a repeated `isnull`, so a column passed twice is checked once. + checkSparkAnswerAndImpl(sql("SELECT addBoth(a, a) FROM t"), dispatched = "addboth" +: guard) + checkSparkAnswerAndImpl( + sql("SELECT addThree(a, b, a) FROM t"), + dispatched = Seq("addthree", "or") ++ guard) + // An argument that cannot be null gets no `isnull`. + checkSparkAnswerAndImpl( + sql("SELECT addBoth(a, 1L) FROM t"), + dispatched = "addboth" +: guard) + } + } + + test("a boolean ScalaUDF's null guard runs in the kernel with false for null (#6704)") { + // In a filter, and in the predicate of a conditional, `ReplaceNullWithFalseInPredicate` + // rewrites the guard to `if (isnull(a)) false else f(knownnotnull(a))`. + spark.udf.register("isPositive", (x: Long) => x > 0) + withTypedCol("BIGINT", "1", "NULL", "-3", "0", "7") { + checkSparkAnswerAndImpl( + sql("SELECT c FROM t WHERE isPositive(c)"), + dispatched = Seq("ispositive", "if", "isnull")) + // The query's own `if` is native, so only the guard's `isnull` is checked. + checkSparkAnswerAndImpl( + sql("SELECT IF(isPositive(c), 'yes', 'no') FROM t"), + dispatched = Seq("ispositive", "isnull")) + } + } + + test("a nondeterministic ScalaUDF under its guard is called only for non-null rows (#6704)") { + // Each result counts the calls made so far in the task, so a call on a null row would shift + // every result after it. + import org.apache.spark.sql.functions.udf + var calls = 0L + val countCalls = udf { (x: Long) => + calls += 1 + x * 1000 + calls + } + spark.udf.register("countCalls", countCalls.asNondeterministic()) + withTypedCol("BIGINT", "1", "NULL", "2", "NULL", "NULL", "3") { + checkSparkAnswerAndImpl( + sql("SELECT countCalls(c) FROM t"), + dispatched = Seq("countcalls", "if", "isnull")) + } + } + + test("a null guard the dispatcher cannot take whole stays native around its UDF (#6704)") { + spark.udf.register("plusOne", (x: Long) => x + 1) + withTypedCol("BIGINT", "1", "NULL", "-3") { + // `canHandle` counts each reference to a column against `maxFields`, and the guard reads + // `c` a second time: the UDF alone counts its output and one input, 2, and the guard 3. At + // 2 the guard is refused while the UDF is not, so the guard has to stay a native `if` + // around the dispatched UDF rather than take the projection to Spark. + withSQLConf("spark.sql.codegen.maxFields" -> "2") { + checkSparkAnswerAndImpl( + sql("SELECT plusOne(c) FROM t"), + native = Seq("if", "isnull"), + dispatched = Seq("plusone")) + } + // With the dispatcher off, the fallback reason is the UDF's own, and none is recorded for + // a guard that the query does not contain. + withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") { + val (_, cometPlan) = checkSparkAnswerAndFallbackReason( + sql("SELECT plusOne(c) FROM t"), + s"plusone: ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=false") + val reasons = new ExtendedExplainInfo().getFallbackReasons(cometPlan) + assert(!reasons.exists(_.startsWith("if:")), s"unexpected fallback reasons: $reasons") + } + } + } + + test("a disabled expression in a null guard keeps the guard native (#6704)") { + // Converting the guard natively checks `spark.comet.expression..enabled` on the UDF and + // on each node of the predicate, the guarded argument's included, and a disabled one takes + // the projection to Spark. Dispatching the guard has to make the same checks. + spark.udf.register("plusOne", (x: Long) => x + 1) + spark.udf.register("addBoth", (a: Long, b: Long) => a + b) + withTable("t") { + sql("CREATE TABLE t (a BIGINT, b BIGINT) USING parquet") + sql("INSERT INTO t VALUES (1, 10), (NULL, 20), (-3, NULL)") + Seq( + "ScalaUDF" -> "SELECT plusOne(a) FROM t", + "IsNull" -> "SELECT plusOne(a) FROM t", + "Or" -> "SELECT addBoth(a, b) FROM t", + "Multiply" -> "SELECT plusOne(a * 2) FROM t").foreach { case (name, query) => + val key = CometConf.getExprEnabledConfigKey(name) + withSQLConf(key -> "false") { + val (_, cometPlan) = checkSparkAnswerAndFallbackReason(sql(query), s"Set $key=true") + assert( + collect(cometPlan) { case p: CometProjectExec => p }.isEmpty, + s"$query stayed in Comet with $key=false:\n$cometPlan") + } + } + } + } + + test("only the null guard Spark builds goes to the kernel with its UDF (#6704)") { + spark.udf.register("plusOne", (x: Long) => x + 1) + withTable("t") { + sql("CREATE TABLE t (a BIGINT, b BIGINT) USING parquet") + val plan = sql("SELECT plusOne(a), b FROM t").queryExecution.optimizedPlan + val guard = plan.expressions.flatMap(_.collect { case g: If => g }).head + val Seq(a, b) = plan.collectLeaves().head.output + def root(expr: Expression): ExprStructCase = + QueryPlanSerde.exprToProto(expr, Seq(a, b)).get.getExprStructCase + assert(root(guard) == ExprStructCase.JVM_SCALAR_UDF) + // A predicate that checks a column the call does not read, or one besides its arguments. + assert(root(guard.copy(predicate = IsNull(b))) == ExprStructCase.IF) + assert(root(guard.copy(predicate = Or(guard.predicate, IsNull(b)))) == ExprStructCase.IF) + // A branch that is not null. + assert(root(guard.copy(trueValue = Literal(0L))) == ExprStructCase.IF) + // A nondeterministic argument, which the guard and the call would each evaluate. + val random = guard.transformUp { case r: AttributeReference => + If(LessThan(Rand(Literal(1L)), Literal(2.0)), r, Literal(null, LongType)) + } + assert(root(random) == ExprStructCase.IF) + } + } + // Arrow Java ignores ArrowArray.offset on import, so the bridge has to zero a sliced boolean's // offset before handing it over, at the top level and inside a struct. // https://github.com/apache/datafusion-comet/issues/6288 @@ -1304,6 +1473,137 @@ class CometCodegenSuite } } + private val boxedPrimitiveTypes = Seq("bool", "byte", "short", "int", "long", "float", "double") + + /** + * Boxed primitive UDFs for the #6706 tests, where the kernel drops Spark's encoders. Each + * `_str` function names the exact value it receives, null included, so a value converted + * differently from Spark's encoder (a null read as 0, a -0.0 read as 0.0) fails the comparison + * with Spark. Each `_id` function returns the boxed value it receives, which covers the + * result side. `long_tag` and `opt_long` pair a boxed parameter with ones that keep their + * encoders. + */ + private def registerBoxedPrimitiveUdfs(): Unit = { + spark.udf.register("bool_str", (x: java.lang.Boolean) => String.valueOf(x)) + spark.udf.register("byte_str", (x: java.lang.Byte) => String.valueOf(x)) + spark.udf.register("short_str", (x: java.lang.Short) => String.valueOf(x)) + spark.udf.register("int_str", (x: java.lang.Integer) => String.valueOf(x)) + spark.udf.register("long_str", (x: java.lang.Long) => String.valueOf(x)) + // Raw bits, so the sign of a zero or a NaN payload would show. + spark.udf.register( + "float_str", + (x: java.lang.Float) => + if (x == null) "null" else java.lang.Float.floatToRawIntBits(x).toString) + spark.udf.register( + "double_str", + (x: java.lang.Double) => + if (x == null) "null" else java.lang.Double.doubleToRawLongBits(x).toString) + spark.udf.register("bool_id", (x: java.lang.Boolean) => x) + spark.udf.register("byte_id", (x: java.lang.Byte) => x) + spark.udf.register("short_id", (x: java.lang.Short) => x) + spark.udf.register("int_id", (x: java.lang.Integer) => x) + spark.udf.register("long_id", (x: java.lang.Long) => x) + spark.udf.register("float_id", (x: java.lang.Float) => x) + spark.udf.register("double_id", (x: java.lang.Double) => x) + spark.udf.register("long_tag", (x: java.lang.Long, s: String) => s"$x:$s") + spark.udf.register("opt_long", (x: Option[Long]) => x.map(_ + 1).getOrElse(-1L)) + } + + private def withBoxedPrimitiveTable(f: => Unit): Unit = { + withTable("t") { + sql( + "CREATE TABLE t (c_bool BOOLEAN, c_byte TINYINT, c_short SMALLINT, c_int INT, " + + "c_long BIGINT, c_float FLOAT, c_double DOUBLE, s STRING) USING parquet") + sql( + "INSERT INTO t VALUES " + + "(true, CAST(-128 AS TINYINT), CAST(-32768 AS SMALLINT), -2147483648, " + + "-9223372036854775808, CAST('-0.0' AS FLOAT), CAST('-0.0' AS DOUBLE), 'a'), " + + "(false, CAST(127 AS TINYINT), CAST(32767 AS SMALLINT), 2147483647, " + + "9223372036854775807, CAST('NaN' AS FLOAT), CAST('NaN' AS DOUBLE), NULL), " + + "(NULL, NULL, NULL, NULL, NULL, NULL, NULL, 'c'), " + + "(true, CAST(0 AS TINYINT), CAST(1000 AS SMALLINT), 1000, 1000, " + + "CAST('Infinity' AS FLOAT), CAST(0.0 AS DOUBLE), 'd'), " + + "(false, CAST(-1 AS TINYINT), CAST(-1000 AS SMALLINT), -1000, -1000, " + + "CAST(1.5 AS FLOAT), CAST('-Infinity' AS DOUBLE), NULL)") + f + } + } + + test("the kernel drops only boxed primitive encoders from a ScalaUDF (#6706)") { + registerBoxedPrimitiveUdfs() + spark.udf.register("str_id", (x: String) => x) + spark.udf.register("long_in_prim_out", (x: java.lang.Long) => if (x == null) -1L else x + 1L) + spark.udf.register("prim_add_one", (x: Long) => x + 1) + withBoxedPrimitiveTable { + def boundUdf(call: String): ScalaUDF = { + val udf = sql(s"SELECT $call FROM t").queryExecution.optimizedPlan.expressions + .flatMap(_.collect { case u: ScalaUDF => u }) + .head + val attrs = udf.collect { case a: AttributeReference => a }.distinct + BindReferences.bindReference(udf, AttributeSeq(attrs)) + } + // Whether the kernel drops the encoder of each parameter and then of the result. + def dropped(call: String): Seq[Boolean] = { + val udf = boundUdf(call) + val kernelUdf = + CometBatchKernelCodegen.withoutBoxedPrimitiveEncoders(udf).asInstanceOf[ScalaUDF] + (udf.inputEncoders :+ udf.outputEncoder) + .zip(kernelUdf.inputEncoders :+ kernelUdf.outputEncoder) + .map { case (before, after) => before.isDefined && after.isEmpty } + } + boxedPrimitiveTypes.foreach { t => + assert(dropped(s"${t}_id(c_$t)") === Seq(true, true), t) + } + assert(dropped("long_in_prim_out(c_long)") === Seq(true, false)) + assert(dropped("long_tag(c_long, s)") === Seq(true, false, false)) + assert(dropped("prim_add_one(c_long)") === Seq(false, false)) + assert(dropped("opt_long(c_long)") === Seq(false, false)) + assert(dropped("str_id(s)") === Seq(false, false)) + + // Spark's encoder path calls the parameter's converter even for a null input. Without the + // encoder, `ScalaUDF` tests for null itself, which shows the kernel compiles the rewrite. + def kernelSource(call: String, vectorClass: Class[_ <: ValueVector]): String = + CometBatchKernelCodegen + .generateSource( + boundUdf(call), + IndexedSeq(ArrowColumnSpec(vectorClass, nullable = true))) + .body + assert(kernelSource("str_id(s)", classOf[VarCharVector]).contains(".apply(null)")) + assert(!kernelSource("long_id(c_long)", classOf[BigIntVector]).contains(".apply(null)")) + } + } + + test("boxed primitive UDF parameters and results match Spark (#6706)") { + registerBoxedPrimitiveUdfs() + withBoxedPrimitiveTable { + val calls = boxedPrimitiveTypes.flatMap(t => Seq(s"${t}_str(c_$t)", s"${t}_id(c_$t)")) + assertCodegenRan { + checkSparkAnswerAndOperator(sql(s"SELECT ${calls.mkString(", ")} FROM t")) + } + } + } + + test("boxed primitive UDF parameters next to other parameters match Spark (#6706)") { + registerBoxedPrimitiveUdfs() + // Spark guards the call with a null check for the primitive parameter, not the boxed one. + spark.udf.register("long_or", (x: java.lang.Long, y: Long) => if (x == null) y else x + y) + withBoxedPrimitiveTable { + val calls = Seq( + "long_tag(c_long, s)", + "opt_long(c_long)", + "long_or(c_long, c_int)", + // Spark casts the int column to the `java.lang.Long` parameter's type. + "long_str(c_int)", + "long_id(long_id(c_long))", + // Inside a lambda, Spark evaluates the UDF through `eval` rather than generated code, + // with the same converters. + "transform(array(c_int, NULL), x -> int_str(x))") + assertCodegenRan { + checkSparkAnswerAndOperator(sql(s"SELECT ${calls.mkString(", ")} FROM t")) + } + } + } + test("ScalaUDF returning a different type than its input") { // String -> Int output transition. Identity-loop above keeps input == output. This asserts // the writer can switch types per the UDF's declared return. @@ -1490,11 +1790,10 @@ class CometCodegenSuite // `init(int partitionIndex)`, called once per kernel allocation. Spark seeds // `XORShiftRandom(seed + partitionIndex)` per partition, so different partitions produce // different sequences for the same seed. Matching Spark across partitions requires the - // kernel to see the real partition index, which the dispatcher derives from - // `TaskContext.get().partitionId()`, live on this path thanks to the bridge-level - // TaskContext propagation. Composing with a ScalaUDF (identity on Double here) forces the - // tree through codegen dispatch so the Rand evaluation runs inside our kernel's init - // rather than via Spark's normal codegen. + // kernel to see the real partition index, which the native plan passes through the bridge. + // Composing with a ScalaUDF (identity on Double here) forces the tree through codegen + // dispatch so the Rand evaluation runs inside our kernel's init rather than via Spark's + // normal codegen. spark.udf.register("dblId", (d: Double) => d) val df = spark .range(0, 1024, 1, numPartitions = 4) @@ -1502,6 +1801,71 @@ class CometCodegenSuite checkSparkAnswerAndOperator(df) } + test("scalar subquery inside a dispatched expression falls back to Spark") { + assume(isSpark40Plus, "collations are Spark 4.0+") + withTable("t") { + sql("CREATE TABLE t (c1 STRING) USING parquet") + sql("INSERT INTO t VALUES ('a')") + checkSparkAnswer(sql("SELECT COUNT(CAST((SELECT c1 FROM t) AS STRING COLLATE UTF8_LCASE))")) + checkSparkAnswer(sql("SELECT CAST((SELECT c1 FROM t) AS STRING COLLATE UTF8_LCASE) = 'A'")) + checkSparkAnswer( + sql("SELECT COUNT(CAST((SELECT c1 FROM t) AS STRING COLLATE UTF8_LCASE)) FROM t")) + } + } + + test("dispatched kernels see the index of the partition the native plan computes") { + // Under a union, a coalesce and a cartesian product the native plan computes a partition + // whose index differs from `TaskContext.partitionId()`, and the coalesce runs two plans in + // one task. See https://github.com/apache/datafusion-comet/issues/6570. + val df = spark + .range(0, 8, 1, numPartitions = 2) + .selectExpr("id", "map(1, spark_partition_id()) AS m", "round(rand(42), 6) AS r") + Seq(df.union(df), df.coalesce(1)).foreach { q => + assertCodegenRan(checkSparkAnswerAndOperator(q)) + } + withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + // `CartesianProductExec` stays in Spark, over native children. + assertCodegenRan(checkSparkAnswer(df.crossJoin(spark.range(0, 4, 1, 2)))) + } + if (isSpark40Plus) { + // Spark 4 lowers `make_valid_utf8` to an `Invoke` whose `deterministic` skips its target + // object, so the partition-seeded input sits under a root that reports deterministic. + val invoke = spark + .range(0, 8, 1, numPartitions = 2) + .selectExpr( + "id", + "make_valid_utf8(cast(spark_partition_id() AS STRING)) AS p", + "make_valid_utf8(cast(round(rand(42), 6) AS STRING)) AS r") + assertCodegenRan(checkSparkAnswerAndOperator(invoke.coalesce(1))) + } + } + + test("plans of a coalesce share deterministic kernels and drop their own when they close") { + // The coalesce runs four native plans in one task. The deterministic `round` compiles once + // for all of them and the nondeterministic one once per plan. Each plan's kernel is dropped + // when the plan closes, so once the rows are read only the shared one is left. + val df = spark + .range(0, 16, 1, numPartitions = 4) + .selectExpr("round(id / 3, 2) AS d", "round(rand(42), 6) AS r") + .coalesce(1) + CometScalaUDFCodegen.resetStats() + val cached = df.queryExecution.toRdd + .mapPartitions { rows => + rows.foreach(_ => ()) + val dispatcher = CometUdfBridge.instanceFor( + TaskContext.get().taskAttemptId(), + classOf[CometScalaUDFCodegen].getName) + Iterator(Option(dispatcher).map(_.asInstanceOf[CometScalaUDFCodegen].cachedPlanIds)) + } + .collect() + .toSeq + assert( + cached == Seq(Some(List(CometScalaUDFCodegen.NoPlan))), + s"expected only the shared kernel to outlive the plans, got plan ids $cached") + val stats = CometScalaUDFCodegen.stats() + assert(stats.compileCount == 5, s"expected 1 shared and 4 per-plan compiles, got $stats") + } + test("ScalaUDF composed with reused scalar subquery across projection and filter") { // The same scalar subquery appears in two sites: the projection (which the dispatcher // compiles into a fused kernel) and the filter (a separate operator). Each site holds its @@ -2133,6 +2497,7 @@ class CometCodegenSuite // because Spark 4.1 still rejects TIME columns in file-based data sources, so no SQL query // can produce a TIME input today. val timeVec = new TimeNanoVector("tm", CometArrowAllocator) + val digestVec = new VarBinaryVector("digest", CometArrowAllocator) val exprVec = new VarBinaryVector("expr", CometArrowAllocator) var out: ValueVector = null try { @@ -2146,21 +2511,72 @@ class CometCodegenSuite val serialized = SparkEnv.get.closureSerializer.newInstance().serialize(expr) val bytes = new Array[Byte](serialized.remaining()) serialized.get(bytes) + digestVec.allocateNew() + digestVec.setSafe(0, CometScalaUDFCodegen.digest(bytes)) + digestVec.setValueCount(1) exprVec.allocateNew() exprVec.setSafe(0, bytes) exprVec.setValueCount(1) - out = new CometScalaUDFCodegen().evaluate(Array(exprVec, timeVec), 2) + out = new CometScalaUDFCodegen().evaluate(Array(digestVec, exprVec, timeVec), 2) val comet = CometVector.getVector(out.asInstanceOf[FieldVector], null) assert(comet.getLong(0) === 45296000000000L) assert(comet.isNullAt(1)) } finally { if (out != null) out.close() + digestVec.close() exprVec.close() timeVec.close() } } + test("dispatcher finds a compiled kernel by the expression digest alone (#6705)") { + // The serde ships a digest of the serialized expression at arg 0 and the bytes at arg 1, and + // the dispatcher reads the bytes only to compile on a cache miss. The second call passes a + // null at arg 1, so it succeeds only if the digest finds the kernel the first call compiled. + val expr = Add(BoundReference(0, LongType, nullable = true), Literal(1L)) + val serialized = SparkEnv.get.closureSerializer.newInstance().serialize(expr) + val bytes = new Array[Byte](serialized.remaining()) + serialized.get(bytes) + + def binaryScalar(name: String, value: Array[Byte]): VarBinaryVector = { + val v = new VarBinaryVector(name, CometArrowAllocator) + v.allocateNew() + if (value == null) v.setNull(0) else v.setSafe(0, value) + v.setValueCount(1) + v + } + val digestVec = binaryScalar("digest", CometScalaUDFCodegen.digest(bytes)) + val exprVec = binaryScalar("expr", bytes) + val nullExprVec = binaryScalar("expr", null) + val input = new BigIntVector("x", CometArrowAllocator) + try { + input.allocateNew(2) + input.set(0, 41L) + input.setNull(1) + input.setValueCount(2) + + val dispatcher = new CometScalaUDFCodegen() + Seq(exprVec, nullExprVec).foreach { arg1 => + val out = + dispatcher.evaluate(Array(digestVec, arg1, input), 2).asInstanceOf[BigIntVector] + try { + assert(out.get(0) === 42L) + assert(out.isNull(1)) + } finally out.close() + } + + // The layout before the digest, with the serialized expression at arg 0, is refused + // rather than taken for a digest. + val e = intercept[IllegalArgumentException] { + dispatcher.evaluate(Array(exprVec, input), 2) + } + assert(e.getMessage.contains("expression digest"), e.getMessage) + } finally { + Seq(digestVec, exprVec, nullExprVec, input).foreach(_.close()) + } + } + // Runtime coverage for nullable nested `getStruct` / `getArray` / `getMap` element reads is // exercised through HOFs in `CometCodegenHOFSuite`. Static emitter assertions live in // `CometCodegenSourceSuite`. diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala index 098b8dd3877..6e1d77b3619 100644 --- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala @@ -242,11 +242,10 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } - // Floats nested in array and struct sort keys are normalized natively as well, but strict mode - // declines a key whose type can hold a null element or field: Spark orders that null below - // every value whatever the key's null order, and the native sort places it by the null order - // (#6476). Every field the generator makes is nullable, so these sorts fall back. A unique `id` - // sorted last makes the ordering total, as above. + // Floats nested in array and struct sort keys are normalized natively, in strict mode too. A key + // whose type can hold a null element or field sorts natively under the default null order, which + // agrees with Spark's. Every field the generator makes is nullable. A unique `id` sorted last + // makes the ordering total, as above. private def checkStrictNestedFloatingPointSort(schema: StructType): Unit = { val df = FuzzDataGenerator.generateDataFrame( new Random(42), @@ -260,21 +259,15 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { df.withColumn("id", monotonically_increasing_id()).write.parquet(path) spark.read.parquet(path).createOrReplaceTempView("tbl") - withSQLConf( - CometConf.getExprAllowIncompatConfigKey("SortOrder") -> "false", - CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "true") { - checkSparkAnswerAndFallbackReason( - sql("select * from tbl order by 1, 2, 3"), - "can hold a null element or field") - } - - // The default null order agrees with Spark's, so opting in keeps the sort native and right. - withSQLConf( - CometConf.getExprAllowIncompatConfigKey("SortOrder") -> "true", - CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "true") { - checkSparkAnswerAndOperator( - sql("select * from tbl order by 1, 2, 3"), - Seq(classOf[CometSortExec])) + // The sort stays native whether or not incompatible expressions are allowed. + Seq("false", "true").foreach { allowIncompat => + withSQLConf( + CometConf.getExprAllowIncompatConfigKey("SortOrder") -> allowIncompat, + CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "true") { + checkSparkAnswerAndOperator( + sql("select * from tbl order by 1, 2, 3"), + Seq(classOf[CometSortExec])) + } } } } @@ -301,8 +294,9 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { test("strict floating point: nested sort keeps NaN payloads and zero signs unchanged") { // As in the scalar test, a local relation keeps the raw bits that Parquet would canonicalize. // `d` is a primitive `Double`, so it is not nullable, and neither are the element of - // `array(d)` and the field of `named_struct('v', d)`: these keys cannot hold a null, and - // strict mode admits them. + // `array(d)` and the field of `named_struct('v', d)`. The `IF` key's element is null on row + // 4, and strict mode admits it too, because its default null order sorts a null element first, + // where Spark does. val negNan = java.lang.Double.longBitsToDouble(0xfff8000000000002L) val posNan = java.lang.Double.longBitsToDouble(0x7ff8000000000002L) val rows = Seq((0, negNan), (1, posNan), (2, -0.0d), (3, 0.0d), (4, 1.0d)) @@ -315,9 +309,14 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { withSQLConf( CometConf.getExprAllowIncompatConfigKey("SortOrder") -> "false", CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "true") { - for ((key, value) <- Seq[(String, Row => Double)]( - "array(d)" -> (_.getSeq[Double](1).head), - "named_struct('v', d)" -> (_.getStruct(1).getDouble(0)))) { + // The zeros are peers, and so are the two NaNs, which sort last. + for ((key, value, order) <- Seq[(String, Row => Option[Double], Seq[Int])]( + ("array(d)", r => Some(r.getSeq[Double](1).head), Seq(2, 3, 4, 0, 1)), + ("named_struct('v', d)", r => Some(r.getStruct(1).getDouble(0)), Seq(2, 3, 4, 0, 1)), + ( + "array(IF(id = 4, CAST(NULL AS DOUBLE), d))", + r => Option(r.getSeq[java.lang.Double](1).head).map(_.doubleValue), + Seq(4, 2, 3, 0, 1)))) { val query = s"SELECT id, $key AS k FROM strict_fp_nested_bits ORDER BY k, id" checkSparkAnswerAndOperator( sql(query), @@ -326,12 +325,16 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { val actual = sql(query).collect().toSeq actual.foreach { row => - assert( - java.lang.Double.doubleToRawLongBits(value(row)) == expected(row.getInt(0)), - s"row ${row.getInt(0)} had its floating-point bits rewritten by the sort on $key") + val id = row.getInt(0) + value(row) match { + case Some(v) => + assert( + java.lang.Double.doubleToRawLongBits(v) == expected(id), + s"row $id had its floating-point bits rewritten by the sort on $key") + case None => assert(id == 4, s"row $id lost its value in the sort on $key") + } } - // The zeros are peers, and so are the two NaNs, which sort last. - assert(actual.map(_.getInt(0)) == Seq(2, 3, 4, 0, 1), s"sort on $key") + assert(actual.map(_.getInt(0)) == order, s"sort on $key") } } } @@ -565,6 +568,30 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("ANSI parse_url of an invalid URL raises a Spark error") { + withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") { + withTable("parse_url_invalid") { + sql("CREATE TABLE parse_url_invalid (u STRING) USING PARQUET") + sql("INSERT INTO parse_url_invalid VALUES ('inva lid://user:pass@host/file')") + + val df = sql("SELECT parse_url(u, 'HOST') FROM parse_url_invalid") + checkCometOperators(stripAQEPlan(df.queryExecution.executedPlan)) + checkSparkAnswerMaybeThrows(df) match { + case (Some(sparkException), Some(cometException)) => + def sparkError(error: Throwable): SparkThrowable = + causeChain(error) + .collectFirst { case error: SparkThrowable => error } + .getOrElse(fail(s"Expected a SparkThrowable, got $error")) + val expected = sparkError(sparkException) + val actual = sparkError(cometException) + assert(actual.getErrorClass == expected.getErrorClass) + assert(actual.getSqlState == expected.getSqlState) + case errors => fail(s"Expected Spark and Comet invalid URL errors, got $errors") + } + } + } + } + test("ANSI decimal divide by zero raises a Spark error") { withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") { withTable("decimal_div_zero") { diff --git a/spark/src/test/scala/org/apache/comet/CometFloatSemanticsSuite.scala b/spark/src/test/scala/org/apache/comet/CometFloatSemanticsSuite.scala index 46683611d50..442cb7c41a7 100644 --- a/spark/src/test/scala/org/apache/comet/CometFloatSemanticsSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometFloatSemanticsSuite.scala @@ -333,26 +333,6 @@ object CometFloatSemanticsSuite { c.group == group && contexts.contains(c.context) val knownGaps: Seq[KnownGap] = Seq( - KnownGap( - issue(6385), - "A comparison outside Project and Filter compares raw Arrow values, where a sign-bit NaN " + - "sorts below every other value. Equi-join conditions already match.", - c => - in("comparison", "aggregate argument", "aggregate FILTER", "sort key", "Generate")(c) || - (c.group == "comparison" && c.context.endsWith("join condition") && - !Set("=", "<=>").contains(c.variant))), - KnownGap( - issue(6385), - "A data filter pushed into the Parquet reader compares a stored column with a sign-bit " + - "NaN literal raw, and drops rows that Spark keeps.", - c => in("literal comparison", "scan filter")(c) && c.variant != "!="), - KnownGap( - issue(6157), - "Ordering and null-safe comparisons of arrays and structs compare float leaves raw.", - c => - c.group == "comparison" && - (c.context.startsWith("array operands") || c.context.startsWith("struct operands")) && - !Set("=", "!=").contains(c.variant)), KnownGap( issue(5312), "collect_set before Spark 4.2 treats -0.0 and 0.0 as one value and NaNs as distinct.", diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergRewriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergRewriteActionSuite.scala index 506e9370fc7..0de8ab12c23 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergRewriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergRewriteActionSuite.scala @@ -373,6 +373,7 @@ class CometIcebergRewriteActionSuite extends CometTestBase with CometIcebergTest case _ => } } + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) spark.sparkContext.addSparkListener(listener) try { body diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergTestBase.scala b/spark/src/test/scala/org/apache/comet/CometIcebergTestBase.scala index 4fe2803db73..ba34b0f951b 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergTestBase.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergTestBase.scala @@ -237,6 +237,9 @@ trait CometIcebergTestBase { this: CometTestBase => override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = if (includeFailures) captured += qe.executedPlan } + // Events from earlier queries may still be queued; drain them so they do not reach the + // listener. + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) spark.listenerManager.register(listener) try { action @@ -259,6 +262,7 @@ trait CometIcebergTestBase { this: CometTestBase => override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = captured += qe.executedPlan } + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) spark.listenerManager.register(listener) try { val error = diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index e13a2ec21ad..36a7a235246 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -38,7 +38,7 @@ import org.json4s.jackson.JsonMethods.parse import org.apache.hadoop.fs.{Path => HadoopPath} import org.apache.parquet.hadoop.ParquetFileReader import org.apache.parquet.hadoop.util.HadoopInputFile -import org.apache.spark.{CometListenerBusUtils, SparkConf, SparkException, Success, TaskContext} +import org.apache.spark.{CometListenerBusUtils, SparkConf, SparkException, SparkThrowable, Success, TaskContext} import org.apache.spark.rdd.RDD import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} import org.apache.spark.sql.CometTestBase @@ -51,12 +51,13 @@ import org.apache.spark.sql.connector.catalog.InMemoryTableCatalog import org.apache.spark.sql.connector.write.{BatchWrite, DataWriterFactory, PhysicalWriteInfo, Write, WriterCommitMessage} import org.apache.spark.sql.execution.{ColumnarToRowTransition, LeafExecNode, SparkPlan} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.streaming.Trigger import org.apache.spark.sql.types.{DoubleType, IntegerType, StringType, StructField, StructType} -import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark41Plus} -import org.apache.comet.iceberg.{IcebergReflection, IcebergWriteReportListener} +import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark41Plus, isSpark42Plus} +import org.apache.comet.iceberg.{IcebergReflection, IcebergWriteReportListener, PositionDeltaWrite} import org.apache.comet.serde.Unsupported import org.apache.comet.serde.operator.CometIcebergNativeWrite @@ -743,6 +744,134 @@ class CometIcebergWriteActionSuite } } + Seq("DELETE", "UPDATE", "MERGE").foreach { command => + test(s"WriteDelta $command is recognized and stays on Iceberg JVM DeltaWriter") { + assume(icebergAvailable, "Iceberg not available in classpath") + assume(isSpark35Plus, "WriteDelta interception starts with Spark 3.5") + withIcebergCatalog { warehouseDir => + val suffix = command.toLowerCase(java.util.Locale.ROOT) + val cometTable = s"write_delta_${suffix}_comet" + val jvmTable = s"write_delta_${suffix}_jvm" + val properties = + "'format-version'='2', 'write.delete.mode'='merge-on-read', " + + "'write.update.mode'='merge-on-read', 'write.merge.mode'='merge-on-read'" + + Seq(cometTable, jvmTable).foreach { table => + createTable( + warehouseDir, + table, + partitionSpec = "PARTITIONED BY (region)", + properties = Some(properties)) + withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "false") { + // Keep ids 1 and 2 in the same Iceberg data file so DELETE id = 1 cannot be + // satisfied as a metadata/file delete and must exercise the WriteDelta path. + coalesceInsert(table, Seq((1, "a", 10.0), (2, "a", 20.0), (3, "c", 30.0))) + } + } + + def run(table: String): Unit = command match { + case "DELETE" => + spark.sql(s"DELETE FROM $catalog.$ns.$table WHERE id = 1") + case "UPDATE" => + spark.sql( + s"UPDATE $catalog.$ns.$table SET region = 'updated', amount = amount + 5 " + + "WHERE id = 2") + case "MERGE" => + spark.sql(s""" + |MERGE INTO $catalog.$ns.$table t + |USING ( + | SELECT 2 AS id, 'updated' AS region, 200.0 AS amount + | UNION ALL + | SELECT 4 AS id, 'new' AS region, 40.0 AS amount + |) s + |ON t.id = s.id + |WHEN MATCHED THEN UPDATE SET t.region = s.region, t.amount = s.amount + |WHEN NOT MATCHED THEN + | INSERT (id, region, amount) VALUES (s.id, s.region, s.amount) + |""".stripMargin) + } + + val jvmBefore = countSnapshots(jvmTable) + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + run(jvmTable) + } + assert(countSnapshots(jvmTable) - jvmBefore == 1L) + + val snapshot = withNativeEnabled { + captureWrite(cometTable) { + run(cometTable) + } + } + assert(snapshot.snapshotDelta == 1L, s"unexpected snapshot count: $snapshot") + + val deltaWrites = snapshot.plans.flatMap { plan => + collectWithSubqueries(plan) { + case writer: IcebergWriteExec if writer.dispatch.isInstanceOf[PositionDeltaWrite] => + writer + } + } + assert( + deltaWrites.nonEmpty, + s"expected PositionDeltaWrite dispatch. Plans:\n${snapshot.plans.mkString("\n--\n")}") + assert( + deltaWrites.forall { writer => + IcebergReflection + .getOuterSparkWrite(writer.batchWrite) + .flatMap(IcebergReflection.getTableFromIcebergWrite) + .isDefined + }, + "expected PositionDelta batch writes to expose their Iceberg table") + + val nativeWrites = snapshot.plans.flatMap { plan => + collectWithSubqueries(plan) { case writer: CometIcebergWriteExec => writer } + } + assert( + nativeWrites.isEmpty, + "WriteDelta must stay on Iceberg's JVM DeltaWriter. Plans:\n" + + snapshot.plans.mkString("\n--\n")) + + def rows(table: String): Seq[Row] = + spark + .sql(s"SELECT id, region, amount FROM $catalog.$ns.$table ORDER BY id") + .collect() + .toSeq + + assert( + rows(cometTable) == rows(jvmTable), + s"$command result differs from the stock Iceberg JVM path") + + val summaryPrefix = command match { + case "DELETE" => "spark.delete." + case "UPDATE" => "spark.update." + case "MERGE" => "spark.merge-into." + } + def rowLevelSummary(table: String): Map[String, String] = + spark + .sql(s"SELECT summary FROM $catalog.$ns.$table.snapshots " + + "ORDER BY committed_at DESC LIMIT 1") + .collect()(0) + .getMap[String, String](0) + .filter { case (key, _) => key.startsWith(summaryPrefix) } + .toMap + + val cometSummary = rowLevelSummary(cometTable) + val jvmSummary = rowLevelSummary(jvmTable) + assert( + cometSummary == jvmSummary, + s"$command snapshot summary differs from the stock Iceberg JVM path: " + + s"Comet=$cometSummary JVM=$jvmSummary") + val summaryExpected = + (command == "MERGE" && isSpark41Plus) || + ((command == "DELETE" || command == "UPDATE") && isSpark42Plus) + if (summaryExpected) { + assert( + cometSummary.nonEmpty, + s"expected Spark row-level summary metrics for $command, got $cometSummary") + } + } + } + } + test("sanity check: Spark's default DELETE path works against a Hadoop catalog") { assume(icebergAvailable, "Iceberg not available in classpath") withIcebergCatalog { warehouseDir => @@ -809,6 +938,114 @@ class CometIcebergWriteActionSuite } } + for (adaptive <- Seq(false, true)) { + test(s"transition-heavy fallback preserves Iceberg writes with AQE=$adaptive") { + assumeNativeAcceleration() + withIcebergCatalog { warehouseDir => + val suffix = if (adaptive) "aqe" else "no_aqe" + val nativeTable = s"transition_native_$suffix" + val fallbackTable = s"transition_fallback_$suffix" + createTable(warehouseDir, nativeTable, partitionSpec = "") + createTable(warehouseDir, fallbackTable, partitionSpec = "") + val values = "(1, 'us-east', 10.5), (2, 'us-west', 20.3), (3, 'eu', 30.7)" + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false") { + assertNativeWriteEngages(nativeTable, Seq(1, 2, 3)) { + spark.sql(s"INSERT INTO $catalog.$ns.$nativeTable VALUES $values") + } + } + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") { + assertNativeWriteDoesNotEngage(fallbackTable, Seq(1, 2, 3)) { + spark.sql(s"INSERT INTO $catalog.$ns.$fallbackTable VALUES $values") + } + } + } + } + } + + // An unpartitioned INSERT never puts an exchange under the write, so it misses #6152: with AQE + // off, unwrapping a transition that sits on a shuffle used to keep walking into the map stage. + // The IN-subquery on a partitioned copy-on-write DELETE is the plan that does, and the AQE-off + // case is the one that fails with `ColumnarBatch cannot be cast to InternalRow`. + // `us-west` keeps a row after deleting id 2. An emptied partition makes the rewrite emit + // nothing, and AQE then replaces the write input with an empty LocalTableScan, so the executed + // plan no longer contains the shuffle this test is checking for. + for (adaptive <- Seq(false, true)) { + test(s"transition-heavy fallback preserves partitioned CoW deletes with AQE=$adaptive") { + assumeNativeAcceleration() + withIcebergCatalog { warehouseDir => + val suffix = if (adaptive) "aqe" else "no_aqe" + val nativeTable = s"transition_cow_native_$suffix" + val fallbackTable = s"transition_cow_fallback_$suffix" + val props = Some("'write.delete.mode'='copy-on-write'") + val spec = "PARTITIONED BY (region)" + createTable(warehouseDir, nativeTable, spec, props) + createTable(warehouseDir, fallbackTable, spec, props) + val seed = + Seq( + (1, "us-east", 10.0), + (2, "us-west", 20.0), + (3, "eu", 30.0), + (4, "us-east", 40.0), + (5, "us-west", 50.0)) + withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "false") { + coalesceInsert(nativeTable, seed) + coalesceInsert(fallbackTable, seed) + } + def delete(table: String): Unit = + spark.sql( + s"DELETE FROM $catalog.$ns.$table WHERE id IN " + + "(SELECT col1 FROM VALUES (2) AS t(col1))") + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false") { + assertNativeWriteEngages(nativeTable, Seq(1, 3, 4, 5)) { + delete(nativeTable) + } + } + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") { + val snapshot = withNativeEnabled { + captureWrite(fallbackTable)(delete(fallbackTable)) + } + assertExactlyOneCommit(snapshot) + val nativeExecs = snapshot.plans.flatMap { plan => + collectWithSubqueries(plan) { case e: CometIcebergWriteExec => e } + } + assert( + nativeExecs.isEmpty, + "transition reversion should restore IcebergWriteExec. Plans:\n" + + snapshot.plans.mkString("\n--\n")) + // AdaptiveSparkPlanExec is a leaf, so only the AQE-aware collect reaches the exchange. + val hasExchange = snapshot.plans.exists { plan => + collect(plan) { case _: ShuffleExchangeLike => true }.nonEmpty + } + assert( + hasExchange, + "the delete must keep a shuffle under the write. Plans:\n" + + snapshot.plans.mkString("\n--\n")) + assertRows(fallbackTable, Seq(1, 3, 4, 5)) + } + + val nativeDirs = partitionDirs(warehouseDir, nativeTable) + val fallbackDirs = partitionDirs(warehouseDir, fallbackTable) + assert( + fallbackDirs == nativeDirs, + s"partition layout fallback=$fallbackDirs native=$nativeDirs") + } + } + } + // What Iceberg's Spark writer stamps for `sort_order_id` on appended files changed across // releases: through 1.10 `SparkWrite$WriterFactory` never wires the table sort order (files // get 0 even on a sorted table); 1.11 added `SparkWriteConf.outputSortOrderId` and stamps the @@ -1381,39 +1618,30 @@ class CometIcebergWriteActionSuite val commits = writeSnapshot.plans.flatMap { plan => collectWithSubqueries(plan) { case c: IcebergCommitExec => c } } - if (mode == "copy-on-write") { - assert( - commits.nonEmpty, - s"expected >= 1 IcebergCommitExec for $mode, got 0. Plans:\n" + - writeSnapshot.plans.mkString("\n--\n")) - } else { - assert( - commits.isEmpty, - s"merge-on-read should stay on Iceberg WriteDelta, got ${commits.size} IcebergCommitExec. Plans:\n" + - writeSnapshot.plans.mkString("\n--\n")) - } - val mergeExecs = writeSnapshot.plans.flatMap { plan => - collectWithSubqueries(plan) { case e: CometMergeRowsExec => e } - } - if (mode == "merge-on-read" && isSpark41Plus) { - val sparkMergeRows = writeSnapshot.plans.flatMap { plan => + assert( + commits.nonEmpty, + s"expected >= 1 IcebergCommitExec for $mode, got 0. Plans:\n" + + writeSnapshot.plans.mkString("\n--\n")) + if (mode == "merge-on-read") { + val deltaWrites = writeSnapshot.plans.flatMap { plan => collectWithSubqueries(plan) { - case e if e.getClass.getSimpleName == "MergeRowsExec" => e + case writer: IcebergWriteExec + if writer.dispatch.isInstanceOf[PositionDeltaWrite] => + writer } } assert( - mergeExecs.isEmpty, - "Spark 4.1+ merge-on-read uses the stock V2 writer and must retain MergeRowsExec") - assert( - sparkMergeRows.nonEmpty, - s"expected Spark MergeRowsExec for Spark 4.1+ $mode. Plans:\n" + - writeSnapshot.plans.mkString("\n--\n")) - } else { - assert( - mergeExecs.nonEmpty, - s"expected CometMergeRowsExec for $mode. Plans:\n" + + deltaWrites.nonEmpty, + s"expected PositionDeltaWrite dispatch for $mode. Plans:\n" + writeSnapshot.plans.mkString("\n--\n")) } + val mergeExecs = writeSnapshot.plans.flatMap { plan => + collectWithSubqueries(plan) { case e: CometMergeRowsExec => e } + } + assert( + mergeExecs.nonEmpty, + s"expected CometMergeRowsExec for $mode. Plans:\n" + + writeSnapshot.plans.mkString("\n--\n")) val nativeWrites = writeSnapshot.plans.flatMap { plan => collectWithSubqueries(plan) { case e: CometIcebergWriteExec => e } @@ -2235,18 +2463,8 @@ class CometIcebergWriteActionSuite "region", "amount") - var nativeFailure: Throwable = null - val nativePlans = withNativeEnabled { - capturePlans(spark, includeFailures = true) { - nativeFailure = intercept[Exception](appendUnclustered(nativeTable, rows)) - } - } - val jvmFailure = intercept[Exception](appendUnclustered(jvmTable, rows)) - assert( - nativePlans.exists(p => - collectWithSubqueries(p) { case e: CometIcebergWriteExec => e }.nonEmpty), - "expected the aborted write to have gone through CometIcebergWriteExec. Plans:\n" + - nativePlans.mkString("\n--\n")) + val (nativeFailure, jvmFailure) = + failedWrites(nativeTable, jvmTable)(appendUnclustered(_, rows)) val nativeError = clusteredWriterError(nativeFailure) val jvmError = clusteredWriterError(jvmFailure) @@ -2293,18 +2511,8 @@ class CometIcebergWriteActionSuite (2, new Timestamp(0L), Array[Byte](0)), (3, new Timestamp(-1500L), Array[Byte](0, 1, -1))).toDF("id", "ts", "bin") - var nativeFailure: Throwable = null - val nativePlans = withNativeEnabled { - capturePlans(spark, includeFailures = true) { - nativeFailure = intercept[Exception](appendUnclustered(nativeTable, rows)) - } - } - val jvmFailure = intercept[Exception](appendUnclustered(jvmTable, rows)) - assert( - nativePlans.exists(p => - collectWithSubqueries(p) { case e: CometIcebergWriteExec => e }.nonEmpty), - "expected the aborted write to have gone through CometIcebergWriteExec. Plans:\n" + - nativePlans.mkString("\n--\n")) + val (nativeFailure, jvmFailure) = + failedWrites(nativeTable, jvmTable)(appendUnclustered(_, rows)) val (_, nativeMessage) = clusteredWriterError(nativeFailure) assert( @@ -2319,6 +2527,51 @@ class CometIcebergWriteActionSuite } } + // https://github.com/apache/datafusion-comet/issues/6234. The native writer pulls its input + // through an Arrow C stream, and Arrow Java hands native only the text of an exception thrown + // while producing a batch, so the user got a CometNativeException instead of Spark's exception. + // The first batch is read on the JVM to derive the stream's schema, which is why the overflow + // has to land past it: row 90000 of a single 100000-row data file, so one task reads it all. + test("native acceleration: an input error past the first batch fails like the JVM writer") { + assumeNativeAcceleration() + withIcebergCatalog { _ => + withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") { + val (nativeTable, jvmTable) = ("overflow_native", "overflow_jvm") + Seq("overflow_src", nativeTable, jvmTable).foreach { table => + spark.sql(s"CREATE TABLE $catalog.$ns.$table (v INT) USING iceberg") + } + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark + .range(0, 100000, 1, 1) + .selectExpr("CAST(id AS INT) AS v") + .writeTo(s"$catalog.$ns.overflow_src") + .append() + } + def insert(table: String): Unit = + spark.sql( + s"INSERT INTO $catalog.$ns.$table " + + s"SELECT v + ${Int.MaxValue - 90000} FROM $catalog.$ns.overflow_src") + + val (nativeFailure, jvmFailure) = failedWrites(nativeTable, jvmTable)(insert) + + def overflow(failure: Throwable): (Int, Class[_], String, String) = { + val chain = causeChain(failure) + val depth = chain.indexWhere(_.isInstanceOf[ArithmeticException]) + assert(depth >= 0, s"no ArithmeticException in the cause chain of $failure") + val error = chain(depth).asInstanceOf[Throwable with SparkThrowable] + (depth, error.getClass, error.getErrorClass, error.getSqlState) + } + assert(overflow(jvmFailure)._3 == "ARITHMETIC_OVERFLOW", s"JVM writer: $jvmFailure") + assert( + overflow(nativeFailure) == overflow(jvmFailure), + s"native writer: $nativeFailure\nJVM writer: $jvmFailure") + Seq(nativeTable, jvmTable).foreach { table => + assert(spark.table(s"$catalog.$ns.$table").count() == 0, s"$table committed rows") + } + } + } + } + test("native acceleration: target-file-size rolls on iceberg-java's 1000-row cadence") { assumeNativeAcceleration() withIcebergCatalog { warehouseDir => @@ -2825,6 +3078,74 @@ class CometIcebergWriteActionSuite } } + // The native writer reserves what its open files hold from the task's memory pool. A fanout + // write keeps a file open for every partition, so over enough partitions it outgrows the pool + // and fails its task, which Spark can retry, instead of growing in memory that no budget + // accounts for until the executor is killed. + test("native acceleration: a fanout write that outgrows the memory pool fails its task") { + assumeNativeAcceleration() + withIcebergCatalog { _ => + val session = spark + import session.implicits._ + // One task writes 64 partitions of 1000 random 500-byte strings, a run of rows per + // partition. Each open file holds about 500 KB, so the pool of about 4 MiB (0.002 of the + // suite's 2 GiB off-heap size) is outgrown long before all 64 files are open. + val rows = 64000 + (0 until rows) + .map(i => (i, s"r${i / 1000}", new scala.util.Random(i).alphanumeric.take(500).mkString)) + .toDF("id", "region", "payload") + .coalesce(1) + .createOrReplaceTempView("fanout_oom_src") + Seq("fanout_oom", "fanout_oom_control").foreach { table => + // No distribution, so one task writes every partition. + spark.sql(s""" + CREATE TABLE $catalog.$ns.$table (id INT, region STRING, payload STRING) + USING iceberg + PARTITIONED BY (region) + TBLPROPERTIES ( + 'write.spark.fanout.enabled'='true', + 'write.distribution-mode'='none') + """) + } + + withNativeEnabled { + withSQLConf(CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> "0.002") { + val (failedPlans, error) = captureFailedPlans(spark) { + spark.sql(s"INSERT INTO $catalog.$ns.fanout_oom SELECT * FROM fanout_oom_src") + } + assert( + error.toSeq + .flatMap(exceptionChain) + .exists(t => + Option(t.getMessage) + .exists(_.contains("Additional allocation failed for IcebergWriteExec"))), + s"expected the native writer's out-of-memory error, got $error") + assert( + failedPlans.exists(p => + collectWithSubqueries(p) { case w: CometIcebergWriteExec => w }.nonEmpty), + s"the failed write did not run natively:\n${failedPlans.mkString("\n--\n")}") + } + assert(countSnapshots("fanout_oom") == 0L, "the failed write must not commit") + assert( + parquetFiles(dataDir("fanout_oom")).isEmpty, + s"the failed task left data files behind: ${parquetFiles(dataDir("fanout_oom"))}") + + // The executor outlived the failure, and with the whole pool the same write succeeds. + val controlPlans = capturePlans(spark) { + spark.sql(s"INSERT INTO $catalog.$ns.fanout_oom_control SELECT * FROM fanout_oom_src") + } + assert( + controlPlans.exists(p => + collectWithSubqueries(p) { case w: CometIcebergWriteExec => w }.nonEmpty), + "the control write did not run natively") + assert( + spark.sql(s"SELECT count(*) FROM $catalog.$ns.fanout_oom_control").head().getLong(0) + == rows) + assert(parquetFiles(dataDir("fanout_oom_control")).size == 64) + } + } + } + test("native acceleration: a pre-native handoff failure cleans up task files") { assumeNativeAcceleration() withIcebergCatalog { warehouseDir => @@ -3090,6 +3411,8 @@ class CometIcebergWriteActionSuite override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = if (taskEnd.reason == Success) JobAbortGate.taskFinished() } + // Drain the setup writes' task-end events first, or they can count toward the gate. + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) spark.sparkContext.addSparkListener(listener) try { val run = () => @@ -3646,6 +3969,28 @@ class CometIcebergWriteActionSuite (depth, chain(depth).getMessage) } + /** + * Runs `write` against `nativeTable` with the native writer and against `jvmTable` with the JVM + * writer, expects both to fail, and returns the two failures. The native write has to go + * through [[CometIcebergWriteExec]], so a fallback to the JVM writer cannot pass as parity. + */ + private def failedWrites(nativeTable: String, jvmTable: String)( + write: String => Unit): (Throwable, Throwable) = { + var nativeFailure: Throwable = null + val nativePlans = withNativeEnabled { + capturePlans(spark, includeFailures = true) { + nativeFailure = intercept[Exception](write(nativeTable)) + } + } + val jvmFailure = intercept[Exception](write(jvmTable)) + assert( + nativePlans.exists(p => + collectWithSubqueries(p) { case e: CometIcebergWriteExec => e }.nonEmpty), + "expected the failed write to have gone through CometIcebergWriteExec. Plans:\n" + + nativePlans.mkString("\n--\n")) + (nativeFailure, jvmFailure) + } + private def captureWrite(tableName: String)(action: => Unit): WriteSnapshot = { val before = countSnapshots(tableName) val plans = capturePlans(spark)(action) @@ -3863,6 +4208,16 @@ class CometIcebergWriteActionSuite "report_orc", partitionSpec = "", properties = Some("'write.format.default'='orc'")) + if (isSpark35Plus) { + createTable( + warehouseDir, + "report_delta", + partitionSpec = "PARTITIONED BY (region)", + properties = Some("'format-version'='2', 'write.delete.mode'='merge-on-read'")) + withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "false") { + coalesceInsert("report_delta", Seq((1, "us-east", 1.5), (2, "us-east", 2.5))) + } + } val writes = reportedWrites { withNativeEnabled { @@ -3872,14 +4227,19 @@ class CometIcebergWriteActionSuite .sql(s"INSERT INTO $catalog.$ns.report_parquet VALUES (1, 'us-east', 1.5)") .collect() spark.sql(s"INSERT INTO $catalog.$ns.report_orc VALUES (2, 'eu', 2.5)") + if (isSpark35Plus) { + spark.sql(s"DELETE FROM $catalog.$ns.report_delta WHERE id = 1") + } } withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "false") { spark.sql(s"INSERT INTO $catalog.$ns.report_parquet VALUES (3, 'eu', 3.5)") } } - assert(writes.map(_.writer) == Seq("native", "jvm", "spark"), writes.mkString("\n")) + val expectedWriters = + Seq("native", "jvm") ++ (if (isSpark35Plus) Seq("jvm") else Nil) ++ Seq("spark") + assert(writes.map(_.writer) == expectedWriters, writes.mkString("\n")) assert(writes(1).reasons.exists(_.contains("only parquet"))) - assert(writes(2).node == "AppendData") + assert(writes.last.node == "AppendData") assert(writes.forall(!_.failed)) } } @@ -3951,6 +4311,8 @@ class CometIcebergWriteActionSuite val listener = new IcebergWriteReportListener( new SparkConf() .set(CometConf.COMET_ICEBERG_WRITE_REPORT_DIR.key, reportDir.getAbsolutePath)) + // Drain the setup writes' events first, or they can reach the listener and be reported. + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) spark.listenerManager.register(listener) try { action diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala index da60c14ff0b..7caffd0e7ee 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala @@ -39,7 +39,7 @@ import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference} import org.apache.spark.sql.comet.{CometIcebergWriteExec, CometSparkToColumnarExec, IcebergWriteExec} import org.apache.spark.sql.execution.{ApplyColumnarRulesAndInsertTransitions, ColumnarToRowExec, CommandExecutionMode, LeafExecNode, SparkPlan} -import org.apache.spark.sql.types.IntegerType +import org.apache.spark.sql.types.{BinaryType, IntegerType} import org.apache.spark.sql.vectorized.ColumnarBatch import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} @@ -545,6 +545,20 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } } + test("fall-back: mixed-case data location scheme (native opens the location verbatim)") { + // OpenDAL strips the scheme prefix from a path case-sensitively, so `S3://` cannot be opened + // natively even though `s3://` can. The gate must match the scheme verbatim and decline it. + withDetectionCatalog { dir => + createTable( + dir, + "mixed_case_scheme", + partitionSpec = "", + properties = + Some("'write.data.path'='S3://nonexistent-bucket/iceberg/db/mixed_case_scheme'")) + assertUnsupportedContainsAllowingWriteFailure("mixed_case_scheme", "storage scheme", "S3") + } + } + test("fall-back: hostless hdfs:/ data location is read as hdfs, not file") { // Hadoop normalises `hdfs:///p` to `hdfs:/p`; with no `://` the gate used to call it `file`. withDetectionCatalog { dir => @@ -758,6 +772,23 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes } } + test("cluster-wide S3A delete, read and committer defaults do not block a native write") { + val hadoopConf = new Configuration(false) + Seq( + "fs.s3a.bulk.delete.page.size" -> "1000", + "fs.s3a.experimental.input.fadvise" -> "random", + "fs.s3a.committer.magic.enabled" -> "true", + "fs.s3a.committer.name" -> "magic", + "fs.s3a.committer.threads" -> "8").foreach { case (k, v) => hadoopConf.set(k, v) } + assert( + CometIcebergNativeWrite.unsupportedHadoopS3Settings(hadoopConf, Some("target")).isEmpty) + + hadoopConf.set("fs.s3a.encryption.algorithm", "SSE-KMS") + assert( + CometIcebergNativeWrite.unsupportedHadoopS3Settings(hadoopConf, Some("target")) == Seq( + "fs.s3a.encryption.algorithm")) + } + test("Hadoop S3A built-in defaults are ignored but custom resources are effective") { def xml(key: String, value: String): java.io.ByteArrayInputStream = new java.io.ByteArrayInputStream(s""" @@ -1851,9 +1882,15 @@ class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTes * directly. */ private def writeChildAfterTransitionRules(source: SparkPlan): SparkPlan = { + val child = CometSparkToColumnarExec(source) + val output = Seq( + AttributeReference(IcebergWriteExec.CommitMessageColumn, BinaryType, nullable = false)()) + val originalPlan = IcebergWriteExec(null, output, child) val write = CometIcebergWriteExec( Operator.newBuilder().build(), - CometSparkToColumnarExec(source), + originalPlan, + child, + output, batchWrite = null, table = null, partitionSpecId = 0) diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index 98ef07b96fb..da07c05f673 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -442,11 +442,11 @@ class CometMapExpressionSuite extends CometTestBase { } } - // `map_contains_key` lowers to `array_contains(map_keys(...), key)`. `CometArrayContains` reports - // Incompatible for floating-point element types, because native `array_contains` compares - // -0.0/+0.0 and NaN bitwise unlike Spark. So under the default config it codegen-dispatches to - // Spark and the query stays native with Spark-exact results and no fallback. Forcing the native - // kernel with `ArrayContains.allowIncompatible=true` makes it serialize its children instead, so + // `map_contains_key` lowers to `array_contains(map_keys(...), key)`. The key here is a cast, + // which `CometArrayContains` keeps on the codegen dispatcher over a nullable array, because Spark + // skips the value for a null array while the native kernel evaluates it first. So under the + // default config the query stays native with Spark-exact results and no fallback. Forcing the + // native kernel with `ArrayContains.allowIncompatible=true` makes it serialize its children, so // the folded double-keyed map literal reaches the literal-expansion path. That map is nested // inside the INT-keyed outer map, past the outer `element_at` key guard which only sees the INT // key. The nested-key walk in the expansion (`mapKeyTypesExpandable`) is what has to decline the @@ -455,8 +455,8 @@ class CometMapExpressionSuite extends CometTestBase { test("map_contains_key over nested floating-point map keys falls back (multirow)") { withParquetTable((1 until 4).map(i => (i, i.toLong)), "tbl") { val negZero = "CAST(concat('-', CAST(_1 - 1 AS STRING), '.0') AS DOUBLE)" - // allowIncompatible flips array_contains from codegen dispatch to the native kernel, so the - // nested-map literal is actually serialized and the expansion guard runs. + // allowIncompatible lifts array_contains' eager-evaluation guard, so the nested-map literal + // is actually serialized and the expansion guard runs. withSQLConf(CometConf.getExprAllowIncompatConfigKey(classOf[ArrayContains]) -> "true") { checkSparkAnswerAndFallbackReason( "SELECT _1 AS id, map_contains_key(" + @@ -467,18 +467,24 @@ class CometMapExpressionSuite extends CometTestBase { } // The collated counterpart of the double-key `map_contains_key` case above: a nested - // `UTF8_LCASE`-keyed map would reach `array_contains` with bytewise comparison. Expansion declines - // the folded literal so the case-insensitive lookup stays on Spark. The outer lookup key is the - // dynamic `_1` so the inner map survives as a literal (a constant key would let Spark fold - // `map_keys` into a collated-string array literal instead, which never reaches this guard). + // `UTF8_LCASE`-keyed map would reach `array_contains` with bytewise comparison. Like the + // floating-point case, `CometArrayContains` reports collated element types Incompatible, so the + // default config codegen-dispatches it; `allowIncompatible=true` forces the native kernel so the + // expansion guard runs. Expansion declines the folded literal so the case-insensitive lookup + // stays on Spark. The outer lookup key is the dynamic `_1` so the inner map survives as a literal + // (a constant key would let Spark fold `map_keys` into a collated-string array literal instead, + // which never reaches this guard). test("map_contains_key over nested collated map keys falls back (multirow)") { assume(isSpark40Plus) + val query = "SELECT _1 AS id, map_contains_key(" + + "element_at(map(1, map(CAST('A1' AS STRING COLLATE UTF8_LCASE), 7)), _1), " + + "CAST('a1' AS STRING COLLATE UTF8_LCASE)) AS present FROM tbl" withParquetTable((1 until 4).map(i => (i, i.toLong)), "tbl") { - checkSparkAnswerAndFallbackReason( - "SELECT _1 AS id, map_contains_key(" + - "element_at(map(1, map(CAST('A1' AS STRING COLLATE UTF8_LCASE), 7)), _1), " + - "CAST('a1' AS STRING COLLATE UTF8_LCASE)) AS present FROM tbl", - "Unsupported data type MapType") + withSQLConf(CometConf.getExprAllowIncompatConfigKey(classOf[ArrayContains]) -> "true") { + checkSparkAnswerAndFallbackReason(query, "Unsupported data type MapType") + } + // Under the default config the collated lookup is dispatched and matches Spark. + checkSparkAnswer(query) } } diff --git a/spark/src/test/scala/org/apache/comet/CometNativeUdfSuite.scala b/spark/src/test/scala/org/apache/comet/CometNativeUdfSuite.scala index 9f36d18ff78..4bad30e2c96 100644 --- a/spark/src/test/scala/org/apache/comet/CometNativeUdfSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometNativeUdfSuite.scala @@ -446,6 +446,41 @@ class CometNativeUdfSuite extends CometTestBase { } } + test("Spark evaluating a native UDF while planning fails the query") { + // Spark's optimizer evaluates a projection over local data itself, so this fails even though + // Comet would take every operator. + CometNativeUDF.register(spark, "add_one_c", libPath, Seq(LongType), LongType) + val e = intercept[Exception] { + sql("SELECT add_one_c(x) FROM VALUES (1L), (2L) AS t(x)").collect() + } + assert( + stackTraceContains(e, "Spark evaluated the call itself while planning"), + s"unhelpful error: $e") + } + + test("a global sort on a native UDF's result runs when the sort is on a column") { + // Spark evaluates a global sort's keys on a sample of rows to choose range bounds, so the + // sort has to be on a column holding the result rather than on the call itself. + CometNativeUDF.register(spark, "add_one_c", libPath, Seq(LongType), LongType) + val df = spark.range(0, 20, 1, 4).selectExpr("add_one_c(19 - id) AS y").orderBy("y") + assert(df.collect().map(_.getLong(0)).toSeq == (1L to 20L)) + } + + test("a type Comet cannot carry natively is refused at registration") { + // Earlier tests register `echo_c`, so drop it first to see that the refusal installs nothing. + sql("DROP TEMPORARY FUNCTION IF EXISTS echo_c") + val e = intercept[IllegalArgumentException] { + CometNativeUDF.register( + spark, + "echo_c", + libPath, + Seq(ObjectType(classOf[String])), + LongType) + } + assert(e.getMessage.contains("Comet has no native representation"), e.getMessage) + assert(!spark.catalog.functionExists("echo_c")) + } + // A native UDF call has to reach its serde wherever Comet converts expressions, not just in a // projection. Filters, join conditions, grouping keys and window partitioning each route // through a different Comet operator, and a regression in any one of them would show up only as diff --git a/spark/src/test/scala/org/apache/comet/CometSqlFileTestSuite.scala b/spark/src/test/scala/org/apache/comet/CometSqlFileTestSuite.scala index b1770b72bfa..72ffb23ebaf 100644 --- a/spark/src/test/scala/org/apache/comet/CometSqlFileTestSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometSqlFileTestSuite.scala @@ -21,9 +21,13 @@ package org.apache.comet import java.io.File +import scala.util.{Failure, Success, Try} + import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.catalyst.expressions.{Abs, CreateArray, Literal} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.ArrayType class CometSqlFileTestSuite extends CometTestBase with AdaptiveSparkPlanHelper { @@ -87,15 +91,56 @@ class CometSqlFileTestSuite extends CometTestBase with AdaptiveSparkPlanHelper { // Disable constant folding so that literal expressions are evaluated by Comet's // native engine rather than being folded away by Spark's optimizer at plan time. - private val constantFoldingExcluded = Seq( - "spark.sql.optimizer.excludedRules" -> - "org.apache.spark.sql.catalyst.optimizer.ConstantFolding") + private val constantFoldingRule = "org.apache.spark.sql.catalyst.optimizer.ConstantFolding" // Most SQL fixtures here predate Spark 4 ANSI default and expect non-ANSI semantics // (silent overflow/null on bad input). Individual files can opt in via their own // --CONFIG line, which appears later in the pair list and wins. private val ansiDisabled = Seq(SQLConf.ANSI_ENABLED.key -> "false") + /** Apply file settings and restore every SQLConf change, including mid-file SET and RESET. */ + private def withFixtureConfigs(file: SqlTestFile)(body: => Unit): Unit = { + val conf = spark.sessionState.conf + val original = conf.getAllConfs + val optimizerKey = SQLConf.OPTIMIZER_EXCLUDED_RULES.key + // Matrix settings are appended to configs, so the last value for a key wins as before. + val configuredRules = file.configs.reverse + .collectFirst { case (`optimizerKey`, value) => + value + } + .toSeq + .flatMap(_.split(",")) + val rules = (configuredRules ++ file.excludedRules ++ + (if (file.constantFoldingEnabled) Seq.empty else Seq(constantFoldingRule))) + .map(_.trim) + .filter(_.nonEmpty) + .distinct + require( + !file.constantFoldingEnabled || !rules.contains(constantFoldingRule), + "ConstantFolding: enabled conflicts with an explicit ConstantFolding exclusion") + // Spark silently ignores unknown and non-excludable rules. A fixture must fail instead of + // claiming to test a path that its optimizer settings never reached. + val optimizer = spark.sessionState.optimizer + val knownRules = optimizer.defaultBatches.flatMap(_.rules).map(_.ruleName).toSet + rules.foreach { rule => + require(knownRules.contains(rule), s"Unknown optimizer rule: $rule") + require( + !optimizer.nonExcludableRules.contains(rule), + s"Cannot exclude optimizer rule: $rule") + } + val configs = ansiDisabled ++ file.configs.filterNot(_._1 == optimizerKey) :+ + (optimizerKey -> rules.mkString(",")) + try { + withSQLConf(configs: _*)(body) + } finally { + val current = conf.getAllConfs + (current.keySet -- original.keySet).foreach(conf.unsetConf) + original.foreach { case (key, value) => + if (current.get(key) != Some(value)) conf.setConfString(key, value) + } + } + } + /** * Pin the sentinel-query convention for fixtures that route an `expect_error` through the * codegen dispatcher. See [[ExpectError]] for the failure mode this guards against. @@ -103,9 +148,9 @@ class CometSqlFileTestSuite extends CometTestBase with AdaptiveSparkPlanHelper { private def requireSentinelForCodegenExpectError( relativePath: String, file: SqlTestFile): Unit = { - val codegenFlagOn = file.configs.exists { case (k, v) => - k == CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key && v.equalsIgnoreCase("true") - } + val codegenFlagOn = spark.sessionState.conf.getAllConfs + .get(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key) + .exists(_.equalsIgnoreCase("true")) if (!codegenFlagOn) return val hasExpectError = file.records.exists { case SqlQuery(_, _: ExpectError, _) => true @@ -132,9 +177,7 @@ class CometSqlFileTestSuite extends CometTestBase with AdaptiveSparkPlanHelper { } private def runTestFile(relativePath: String, file: SqlTestFile): Unit = { - requireSentinelForCodegenExpectError(relativePath, file) - val allConfigs = ansiDisabled ++ file.configs ++ constantFoldingExcluded - withSQLConf(allConfigs: _*) { + withFixtureConfigs(file) { withTable(file.tables: _*) { file.records.foreach { case SqlStatement(sql, line) => @@ -167,6 +210,8 @@ class CometSqlFileTestSuite extends CometTestBase with AdaptiveSparkPlanHelper { case Ignore(reason) => logInfo(s"IGNORED query ($reason): $sql") case ExpectError(pattern) => + // Check the effective setting after any mid-file SET or RESET. + requireSentinelForCodegenExpectError(relativePath, file) val (sparkError, cometError) = checkSparkAnswerMaybeThrows(spark.sql(sql)) assert( sparkError.isDefined, @@ -192,39 +237,198 @@ class CometSqlFileTestSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("SQL fixture optimizer directives control the optimized expression") { + val default = SqlFileTestParser.parse(Seq.empty) + withFixtureConfigs(default) { + val plan = spark.sql("SELECT array(1, 2)").queryExecution.optimizedPlan + assert(plan.expressions.exists(_.exists(_.isInstanceOf[CreateArray]))) + val nullPlan = spark.sql("SELECT abs(CAST(NULL AS INT))").queryExecution.optimizedPlan + assert(!nullPlan.expressions.exists(_.exists(_.isInstanceOf[Abs]))) + } + val folded = SqlFileTestParser.parse(Seq("-- ConstantFolding: enabled")) + withFixtureConfigs(folded) { + val plan = spark.sql("SELECT array(1, 2)").queryExecution.optimizedPlan + assert(plan.expressions.exists(_.exists { + case Literal(_, _: ArrayType) => true + case _ => false + })) + assert(!plan.expressions.exists(_.exists(_.isInstanceOf[CreateArray]))) + } + val excluded = SqlFileTestParser.parse( + Seq("-- ExcludeRules: org.apache.spark.sql.catalyst.optimizer.NullPropagation")) + withFixtureConfigs(excluded) { + val plan = spark.sql("SELECT abs(CAST(NULL AS INT))").queryExecution.optimizedPlan + assert(plan.expressions.exists(_.exists(_.isInstanceOf[Abs]))) + assert(spark.conf.get(SQLConf.OPTIMIZER_EXCLUDED_RULES.key).contains(constantFoldingRule)) + } + } + + test("SQL fixture optimizer Config and matrix settings retain precedence") { + val optimizerKey = SQLConf.OPTIMIZER_EXCLUDED_RULES.key + val nullPropagation = "org.apache.spark.sql.catalyst.optimizer.NullPropagation" + val eliminateSorts = "org.apache.spark.sql.catalyst.optimizer.EliminateSorts" + val file = SqlFileTestParser.parse( + Seq( + s"-- Config: $optimizerKey=$nullPropagation", + s"-- ConfigMatrix: $optimizerKey=$eliminateSorts,$nullPropagation", + s"-- ExcludeRules: $nullPropagation, $nullPropagation")) + val original = spark.sessionState.conf.getAllConfs + configMatrix(file.configMatrix).foreach { matrix => + withFixtureConfigs(file.copy(configs = file.configs ++ matrix)) { + val selected = matrix.head._2 + assert( + spark.conf.get(optimizerKey).split(",").toSeq == + Seq(selected, nullPropagation, constantFoldingRule).distinct) + } + assert(spark.sessionState.conf.getAllConfs == original) + } + } + + test("SQL fixture optimizer directives reject ineffective exclusions") { + Seq( + "org.apache.spark.sql.catalyst.optimizer.NoSuchRule" -> "Unknown optimizer rule", + spark.sessionState.optimizer.nonExcludableRules.head -> "Cannot exclude optimizer rule") + .foreach { case (rule, message) => + val file = SqlFileTestParser.parse(Seq(s"-- ExcludeRules: $rule")) + assert( + intercept[IllegalArgumentException](withFixtureConfigs(file) {}).getMessage + .contains(message)) + } + val conflict = SqlFileTestParser.parse( + Seq("-- ConstantFolding: enabled", s"-- ExcludeRules: $constantFoldingRule")) + assert( + intercept[IllegalArgumentException](withFixtureConfigs(conflict) {}).getMessage + .contains("conflicts")) + } + + test("SQL fixture SET and RESET restore existing and previously unset keys") { + val key = "spark.comet.sqlFixture.isolation" + val original = spark.sessionState.conf.getAllConfs + assert(!original.contains(key)) + val file = SqlFileTestParser.parse(Seq("statement", s"SET $key=first")) + runTestFile("set-isolation", file) + assert(spark.sessionState.conf.getAllConfs == original) + withSQLConf(key -> "original") { + val before = spark.sessionState.conf.getAllConfs + runTestFile("set-existing", file) + assert(spark.sessionState.conf.getAllConfs == before) + runTestFile("reset-existing", SqlFileTestParser.parse(Seq("statement", s"RESET $key"))) + assert(spark.sessionState.conf.getAllConfs == before) + runTestFile("reset-all", SqlFileTestParser.parse(Seq("statement", "RESET"))) + assert(spark.sessionState.conf.getAllConfs == before) + } + } + + test("SQL fixture settings are restored after a failing statement or query") { + val key = "spark.comet.sqlFixture.isolation" + Seq("statement", "query").foreach { directive => + val original = spark.sessionState.conf.getAllConfs + val file = SqlFileTestParser.parse( + Seq( + "statement", + s"SET $key=changed", + "", + directive, + "SELECT * FROM missing_sql_fixture_isolation_table")) + intercept[RuntimeException](runTestFile("failure-isolation", file)) + assert(spark.sessionState.conf.getAllConfs == original) + } + } + + test("SQL fixture SET changes take effect between queries") { + val file = SqlFileTestParser.parse( + Seq( + "statement", + "CREATE TABLE sql_fixture_set(v STRING) USING parquet", + "", + "statement", + "INSERT INTO sql_fixture_set VALUES ('invalid'), ('1')", + "", + "query", + "SELECT CAST(v AS INT) FROM sql_fixture_set", + "", + "statement", + "SET spark.sql.ansi.enabled=true", + "", + "query expect_error(CAST_INVALID_INPUT)", + "SELECT CAST(v AS INT) FROM sql_fixture_set")) + val original = spark.sessionState.conf.getAllConfs + runTestFile("set-between-queries", file) + assert(spark.sessionState.conf.getAllConfs == original) + } + + test("SQL fixture SET cannot bypass the dispatcher error sentinel") { + val file = SqlFileTestParser.parse( + Seq( + "-- Config: spark.sql.ansi.enabled=true", + s"-- Config: ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=false", + "statement", + s"SET ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=true", + "", + "query expect_error(CAST_INVALID_INPUT)", + "SELECT CAST('invalid' AS INT)")) + val original = spark.sessionState.conf.getAllConfs + val error = intercept[Exception](runTestFile("set-dispatcher-sentinel", file)) + assert(error.getMessage.contains("missing a non-error")) + assert(spark.sessionState.conf.getAllConfs == original) + } + + test("SQL fixture dispatcher sentinel follows SET before an error query") { + val file = SqlFileTestParser.parse( + Seq( + "-- Config: spark.sql.ansi.enabled=true", + s"-- Config: ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=true", + "statement", + s"SET ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=false", + "", + "query expect_error(CAST_INVALID_INPUT)", + "SELECT CAST('invalid' AS INT)")) + val original = spark.sessionState.conf.getAllConfs + runTestFile("set-disable-dispatcher", file) + assert(spark.sessionState.conf.getAllConfs == original) + } + // Discover and register all .sql test files discoverTestFiles(testResourceDir).foreach { file => val relativePath = testResourceDir.toURI.relativize(file.toURI).getPath - val parsed = SqlFileTestParser.parse(file) - val combinations = configMatrix(parsed.configMatrix) - - // Skip tests that fall outside the file's declared Spark version range. - val skip = skipReason(parsed) - - if (combinations.size <= 1) { - // No matrix or single combination - test(s"sql-file: $relativePath") { - skip match { - case Some(reason) => - logInfo(s"SKIPPED ($reason): $relativePath") - case None => - val effectiveConfigs = parsed.configs ++ combinations.headOption.getOrElse(Seq.empty) - runTestFile(relativePath, parsed.copy(configs = effectiveConfigs)) - } - } - } else { - // Multiple combinations: generate one test per combination - combinations.foreach { matrixConfigs => - val label = matrixConfigs.map { case (k, v) => s"$k=$v" }.mkString(", ") - test(s"sql-file: $relativePath [$label]") { - skip match { - case Some(reason) => - logInfo(s"SKIPPED ($reason): $relativePath") - case None => - runTestFile(relativePath, parsed.copy(configs = parsed.configs ++ matrixConfigs)) + Try(SqlFileTestParser.parse(file)) match { + case Failure(error) => + // A malformed file must not prevent every other fixture from being registered. + test(s"sql-file: $relativePath") { throw error } + case Success(parsed) => + val combinations = configMatrix(parsed.configMatrix) + + // Skip tests that fall outside the file's declared Spark version range. + val skip = skipReason(parsed) + + if (combinations.size <= 1) { + // No matrix or single combination + test(s"sql-file: $relativePath") { + skip match { + case Some(reason) => + logInfo(s"SKIPPED ($reason): $relativePath") + case None => + val effectiveConfigs = + parsed.configs ++ combinations.headOption.getOrElse(Seq.empty) + runTestFile(relativePath, parsed.copy(configs = effectiveConfigs)) + } + } + } else { + // Multiple combinations: generate one test per combination + combinations.foreach { matrixConfigs => + val label = matrixConfigs.map { case (k, v) => s"$k=$v" }.mkString(", ") + test(s"sql-file: $relativePath [$label]") { + skip match { + case Some(reason) => + logInfo(s"SKIPPED ($reason): $relativePath") + case None => + runTestFile( + relativePath, + parsed.copy(configs = parsed.configs ++ matrixConfigs)) + } + } } } - } } } } diff --git a/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala index 2e73ac9c331..e61a57a6418 100644 --- a/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala @@ -41,6 +41,14 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH private val crossTimezones = Seq("UTC", "America/Los_Angeles", "Europe/London", "Asia/Tokyo") + /** + * Dates around 2024 stay inside DataFusion `date_trunc`'s TimestampNanosecond range + * (approximately 1678 through 2262). The fuzz generator's default `baseDate` is year 3333, + * which only exercises Comet's wide-range fallback. + */ + private val dataFusionRangeBaseDate: Long = + new java.text.SimpleDateFormat("yyyy-MM-dd HH:mm:ss").parse("2024-06-15 12:00:00").getTime + private def deepestSparkThrowable(error: Throwable): SparkThrowable with Throwable = causeChain(error) .collect { case e: SparkThrowable with Throwable => e } @@ -80,6 +88,14 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH } test("trunc (TruncDate)") { + checkTruncDate(dataFusionRangeBaseDate) + } + + test("trunc (TruncDate) - dates outside DataFusion nanosecond range") { + checkTruncDate(FuzzDataGenerator.defaultBaseDate) + } + + private def checkTruncDate(baseDate: Long): Unit = { val supportedFormats = CometTruncDate.supportedFormats val unsupportedFormats = Seq("invalid") @@ -88,7 +104,12 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH Seq( StructField("c0", DataTypes.DateType, true), StructField("c1", DataTypes.StringType, true))) - val df = FuzzDataGenerator.generateDataFrame(r, spark, schema, 1000, DataGenOptions()) + val df = FuzzDataGenerator.generateDataFrame( + r, + spark, + schema, + 1000, + DataGenOptions(baseDate = baseDate)) df.createOrReplaceTempView("tbl") @@ -146,7 +167,7 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH val supportedFormats = CometTruncTimestamp.supportedFormats val unsupportedFormats = Seq("invalid") - createTimestampTestData().createOrReplaceTempView("tbl") + createTimestampTestData(dataFusionRangeBaseDate).createOrReplaceTempView("tbl") // TODO test fails with non-UTC timezone // https://github.com/apache/datafusion-comet/issues/2649 @@ -172,7 +193,9 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH val unsupportedFormats = Seq("invalid") withTempDir { path => - createTimestampTestData().write.mode(SaveMode.Overwrite).parquet(path.toString) + createTimestampTestData(dataFusionRangeBaseDate).write + .mode(SaveMode.Overwrite) + .parquet(path.toString) spark.read.parquet(path.toString).createOrReplaceTempView("tbl") // TODO test fails with non-UTC timezone @@ -197,6 +220,15 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH } } + test("date_trunc (TruncTimestamp) - timestamps outside DataFusion nanosecond range") { + createTimestampTestData(FuzzDataGenerator.defaultBaseDate).createOrReplaceTempView("tbl") + withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC") { + for (format <- Seq("year", "week", "hour")) { + checkSparkAnswerAndOperator(s"SELECT c0, date_trunc('$format', c0) from tbl order by c0") + } + } + } + test("date_trunc - non-UTC timezone routes via codegen dispatcher") { // The native date_trunc is Incompatible in non-UTC sessions // (https://github.com/apache/datafusion-comet/issues/2649), so with allowIncompatible=false @@ -241,9 +273,7 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH // for any non-UTC timezone whose DST transitions are 1-hour and whose dates fall within // chrono-tz's precomputed DST horizon (currently ~year 2100). A base date in 2024 keeps the // generator well inside that window. - val baseDate2024 = - new java.text.SimpleDateFormat("yyyy-MM-dd HH:mm:ss").parse("2024-06-15 12:00:00").getTime - createTimestampTestData(baseDate2024).createOrReplaceTempView("tbl") + createTimestampTestData(dataFusionRangeBaseDate).createOrReplaceTempView("tbl") val nonUtcTimezones = Seq("America/New_York", "Europe/London", "Asia/Tokyo") for (tz <- nonUtcTimezones) { @@ -327,14 +357,12 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH // Use a reasonable date range (around year 2024) to avoid chrono-tz DST calculation // issues with far-future dates. The default baseDate is year 3333 which is beyond // the range where chrono-tz can reliably calculate DST transitions. - val reasonableBaseDate = - new java.text.SimpleDateFormat("yyyy-MM-dd HH:mm:ss").parse("2024-06-15 12:00:00").getTime val ntzDF = FuzzDataGenerator.generateDataFrame( r, spark, ntzSchema, 100, - DataGenOptions(baseDate = reasonableBaseDate)) + DataGenOptions(baseDate = dataFusionRangeBaseDate)) ntzDF.createOrReplaceTempView("ntz_tbl") for (tz <- crossTimezones) { withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> tz) { @@ -811,7 +839,7 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> timezone) { checkDays( tsDF.select(col("ts"), getColumnFromExpression(Days(UnresolvedAttribute("ts")))), - tsDF.selectExpr("ts", "unix_date(cast(ts as date))")) + tsDF.selectExpr("ts", "cast(floor(unix_micros(ts) / 86400000000D) as int)")) } } } @@ -857,8 +885,16 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH java.sql.Timestamp.valueOf("2024-06-15 10:30:00"), DataTypes.TimestampType)))), dummyDF.selectExpr( - "unix_date(cast(TIMESTAMP('1970-01-01 00:00:00') as date))", - "unix_date(cast(TIMESTAMP('2024-06-15 10:30:00') as date))")) + "cast(floor(unix_micros(TIMESTAMP('1970-01-01 00:00:00')) / 86400000000D) as int)", + "cast(floor(unix_micros(TIMESTAMP('2024-06-15 10:30:00')) / 86400000000D) as int)")) + + // One microsecond before the epoch. Days are counted in UTC and floored, so this is day -1 + // even though it is already 1970-01-01 in Asia/Tokyo. + withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> "Asia/Tokyo") { + checkDays( + dummyDF.select(getColumnFromExpression(Days(Literal(-1L, DataTypes.TimestampType)))), + dummyDF.selectExpr("-1")) + } // Null handling checkDays( diff --git a/spark/src/test/scala/org/apache/comet/SqlFileTestParser.scala b/spark/src/test/scala/org/apache/comet/SqlFileTestParser.scala index 2b7c573bb19..f7e653c2e13 100644 --- a/spark/src/test/scala/org/apache/comet/SqlFileTestParser.scala +++ b/spark/src/test/scala/org/apache/comet/SqlFileTestParser.scala @@ -104,6 +104,10 @@ case class ExpectError(pattern: String) extends QueryAssertionMode * Optional maximum Spark version this test applies to (e.g. "3.4"). The test is skipped on * newer versions. Useful for paired fixtures where each version range has its own expected * error class or output format. + * @param excludedRules + * Additional optimizer rule names to exclude, alongside the default ConstantFolding exclusion. + * @param constantFoldingEnabled + * Whether to let Spark fold constants before Comet sees the query. */ case class SqlTestFile( configs: Seq[(String, String)], @@ -111,12 +115,18 @@ case class SqlTestFile( records: Seq[SqlTestRecord], tables: Seq[String], minSparkVersion: Option[String] = None, - maxSparkVersion: Option[String] = None) + maxSparkVersion: Option[String] = None, + excludedRules: Seq[String] = Seq.empty, + constantFoldingEnabled: Boolean = false) object SqlFileTestParser { - private val ConfigPattern = """--\s*Config:\s*(.+)=(.+)""".r - private val ConfigMatrixPattern = """--\s*ConfigMatrix:\s*(.+)=(.+)""".r + private val ConfigPattern = """--\s*Config:\s*([^=]+)=(.*)""".r + private val ConfigMatrixPattern = """--\s*ConfigMatrix:\s*([^=]+)=(.*)""".r + private val ExcludeRulesPattern = """--\s*ExcludeRules:\s*(.*)""".r + private val ConstantFoldingPattern = """--\s*ConstantFolding:\s*(.*)""".r + private val ConfigDirectivePattern = + """--\s*(Config|ConfigMatrix|ExcludeRules|ConstantFolding):.*""".r private val MinSparkVersionPattern = """--\s*MinSparkVersion:\s*(.+)""".r private val MaxSparkVersionPattern = """--\s*MaxSparkVersion:\s*(.+)""".r private val CreateTablePattern = """(?i)CREATE\s+TABLE\s+(\w+)""".r.unanchored @@ -125,6 +135,9 @@ object SqlFileTestParser { val source = Source.fromFile(file, "UTF-8") try { parse(source.getLines().toSeq) + } catch { + case e: IllegalArgumentException => + throw new IllegalArgumentException(s"${file.getPath}: ${e.getMessage}", e) } finally { source.close() } @@ -135,6 +148,8 @@ object SqlFileTestParser { var configMatrix = Seq.empty[(String, Seq[String])] var minSparkVersion: Option[String] = None var maxSparkVersion: Option[String] = None + var excludedRules = Seq.empty[String] + var constantFoldingEnabled: Option[Boolean] = None val records = Seq.newBuilder[SqlTestRecord] val tables = Seq.newBuilder[String] @@ -144,13 +159,38 @@ object SqlFileTestParser { line match { case ConfigPattern(key, value) => + require(key.trim.nonEmpty, s"Empty Config key at line ${lineIdx + 1}") configs :+= (key.trim -> value.trim) lineIdx += 1 case ConfigMatrixPattern(key, values) => - configMatrix :+= (key.trim -> values.split(",").map(_.trim).toSeq) + require(key.trim.nonEmpty, s"Empty ConfigMatrix key at line ${lineIdx + 1}") + val choices = values.split(",", -1).map(_.trim).toSeq + require(choices.forall(_.nonEmpty), s"Empty ConfigMatrix value at line ${lineIdx + 1}") + configMatrix :+= (key.trim -> choices) lineIdx += 1 + case ExcludeRulesPattern(rules) => + val names = splitNames(rules) + require(names.nonEmpty, s"Empty ExcludeRules at line ${lineIdx + 1}") + excludedRules ++= names + lineIdx += 1 + + case ConstantFoldingPattern(mode) => + require( + mode == "enabled" || mode == "disabled", + s"Invalid ConstantFolding mode '$mode' at line ${lineIdx + 1}: " + + "expected enabled or disabled") + val enabled = mode == "enabled" + require( + constantFoldingEnabled.forall(_ == enabled), + s"Conflicting ConstantFolding directives at line ${lineIdx + 1}") + constantFoldingEnabled = Some(enabled) + lineIdx += 1 + + case ConfigDirectivePattern(name) => + throw new IllegalArgumentException(s"Malformed $name directive at line ${lineIdx + 1}") + case MinSparkVersionPattern(version) => minSparkVersion = Some(version.trim) lineIdx += 1 @@ -188,7 +228,9 @@ object SqlFileTestParser { records.result(), tables.result(), minSparkVersion, - maxSparkVersion) + maxSparkVersion, + excludedRules.distinct, + constantFoldingEnabled.getOrElse(false)) } private val FallbackPattern = """query\s+expect_fallback\((.+)\)""".r diff --git a/spark/src/test/scala/org/apache/comet/SqlFileTestParserSuite.scala b/spark/src/test/scala/org/apache/comet/SqlFileTestParserSuite.scala index c356c3c37a6..0d107395ba7 100644 --- a/spark/src/test/scala/org/apache/comet/SqlFileTestParserSuite.scala +++ b/spark/src/test/scala/org/apache/comet/SqlFileTestParserSuite.scala @@ -20,6 +20,8 @@ package org.apache.comet import java.io.File +import java.nio.charset.StandardCharsets +import java.nio.file.Files import org.scalatest.funsuite.AnyFunSuite @@ -78,6 +80,83 @@ class SqlFileTestParserSuite extends AnyFunSuite { } } + test("configuration values preserve equals signs and empty Config values") { + val parsed = SqlFileTestParser.parse( + Seq( + "-- Config: fixture.key = MDEyMzQ1Njc4OTAxMjM0NQ==", + "-- Config: fixture.other = a=b=c", + "-- Config: fixture.empty=", + "-- ConfigMatrix: fixture.matrix = first==, second=part=")) + assert( + parsed.configs == Seq( + "fixture.key" -> "MDEyMzQ1Njc4OTAxMjM0NQ==", + "fixture.other" -> "a=b=c", + "fixture.empty" -> "")) + assert(parsed.configMatrix == Seq("fixture.matrix" -> Seq("first==", "second=part="))) + } + + test("optimizer directives preserve defaults and combine repeated exclusions") { + val defaults = SqlFileTestParser.parse(Seq.empty) + assert(!defaults.constantFoldingEnabled) + assert(defaults.excludedRules.isEmpty) + val parsed = SqlFileTestParser.parse( + Seq( + "-- ConstantFolding: enabled", + "-- ExcludeRules: first.Rule, second.Rule,", + "-- ExcludeRules: second.Rule, third.Rule")) + assert(parsed.constantFoldingEnabled) + assert(parsed.excludedRules == Seq("first.Rule", "second.Rule", "third.Rule")) + assert(!SqlFileTestParser.parse(Seq("-- ConstantFolding: disabled")).constantFoldingEnabled) + } + + test("malformed configuration directives fail with their source line") { + Seq( + "-- Config: no_assignment", + "-- Config: =value", + "-- ConfigMatrix: no_assignment", + "-- ConfigMatrix: key=", + "-- ConfigMatrix: key=first,,second", + "-- ConfigMatrix: key=first,", + "-- ExcludeRules:", + "-- ExcludeRules: , ,", + "-- ConstantFolding: true").foreach { directive => + val error = intercept[IllegalArgumentException] { + SqlFileTestParser.parse(Seq("-- comment", directive)) + } + assert(error.getMessage.contains("line 2"), directive) + } + // Prose mentioning a directive is still an ordinary comment unless it includes the colon. + assert( + SqlFileTestParser.parse(Seq("-- ConstantFolding, so literals remain visible")) == + SqlTestFile(Seq.empty, Seq.empty, Seq.empty, Seq.empty)) + } + + test("repeated ConstantFolding directives must agree") { + Seq("enabled", "disabled").foreach { mode => + val directive = s"-- ConstantFolding: $mode" + val parsed = SqlFileTestParser.parse(Seq(directive, directive)) + assert(parsed.constantFoldingEnabled == (mode == "enabled")) + val other = if (mode == "enabled") "disabled" else "enabled" + val error = intercept[IllegalArgumentException] { + SqlFileTestParser.parse(Seq(directive, s"-- ConstantFolding: $other")) + } + assert(error.getMessage.contains("Conflicting ConstantFolding directives at line 2")) + } + } + + test("malformed fixture files report their path and source line") { + val tempDir = Files.createDirectories(new File(System.getProperty("java.io.tmpdir")).toPath) + val path = Files.createTempFile(tempDir, "comet-malformed-fixture-", ".sql") + try { + Files.write(path, "-- comment\n-- Config: no_assignment".getBytes(StandardCharsets.UTF_8)) + val error = intercept[IllegalArgumentException](SqlFileTestParser.parse(path.toFile)) + assert(error.getMessage.contains(path.toString)) + assert(error.getMessage.contains("line 2")) + } finally { + Files.delete(path) + } + } + test("statements and queries preserve SQL, source lines and tables for cleanup") { val lines = Seq( "-- fixture comment", diff --git a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala index c6a6805c72f..a84d4df3e7d 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala @@ -419,6 +419,25 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar } } + test("columnar shuffle on array [timestamp]") { + // The native row-to-Arrow conversion appends an array of 64 or more elements that holds a + // null in one call, which requires the array to carry the column's timezone. + Seq(10, 201).foreach { numPartitions => + withParquetTable( + (0 until 50).map { i => + val timestamps: Seq[java.sql.Timestamp] = (0 until 70).map { j => + if (j == i) null else new java.sql.Timestamp(i * 1000000L + j) + } + (i, timestamps) + }, + "tbl") { + val df = sql("SELECT * FROM tbl").repartition(numPartitions, $"_1") + + checkShuffleAnswer(df, 1) + } + } + } + test("columnar shuffle on nested array") { Seq("false", "true").foreach { _ => Seq(10, 201).foreach { numPartitions => diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExec3_4PlusSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExec3_4PlusSuite.scala index c6924eb91c4..516ece40be0 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExec3_4PlusSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExec3_4PlusSuite.scala @@ -158,6 +158,16 @@ class CometExec3_4PlusSuite extends CometTestBase { } } + test("test NULL inputs for BloomFilterMightContain without a FROM clause") { + assume(!isSpark42Plus, "https://github.com/apache/datafusion-comet/issues/4968") + checkSparkAnswer(""" + |SELECT might_contain(null, null) both_null, + | might_contain(null, 1L) null_bf, + | might_contain((SELECT bloom_filter_agg(cast(id as long)) FROM range(1, 10000)), + | null) null_value + |""".stripMargin) + } + test("test NULL inputs for BloomFilterMightContain") { assume(!isSpark42Plus, "https://github.com/apache/datafusion-comet/issues/4968") val table = "test" diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index d042992c696..478c62f4eda 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -34,6 +34,7 @@ import org.apache.spark.sql.catalyst.{FunctionIdentifier, TableIdentifier} import org.apache.spark.sql.catalyst.catalog.{BucketSpec, CatalogStatistics, CatalogTable} import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression, ExpressionInfo, Hex, Literal} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, BloomFilterAggregate, Final} +import org.apache.spark.sql.catalyst.plans.logical.Expand import org.apache.spark.sql.comet._ import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.connector.catalog.InMemoryTableCatalog @@ -51,7 +52,7 @@ import org.apache.spark.sql.internal.SQLConf.SESSION_LOCAL_TIMEZONE import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.UTF8String -import org.apache.comet.{CometConf, CometExecIterator, ExtendedExplainInfo} +import org.apache.comet.{CometConf, CometExecIterator, CometNativeException, ExtendedExplainInfo} import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus, isSpark41Plus, isSpark42Plus} import org.apache.comet.rules.CometCoalesceShufflePartitions import org.apache.comet.serde.Config.ConfigMap @@ -111,6 +112,51 @@ class CometExecSuite extends CometTestBase { } } + test("native sort spill files are compressed with the default spill codec") { + val numRows = 20000 + val compressibleValue = "native-sort-spill-compression-" * 8 + withTempPath { path => + // A single input file keeps all rows in one task so its sort outgrows the tiny memory + // pool below and must spill. + spark + .createDataFrame((0 until numRows).map(i => (i, compressibleValue))) + .coalesce(1) + .write + .parquet(path.getAbsolutePath) + + withParquetTable(path.getAbsolutePath, "tbl") { + def sortSpilledBytes(codecConf: (String, String)*): Long = { + var spilledBytes = 0L + withSQLConf( + Seq( + CometConf.COMET_BATCH_SIZE.key -> "1024", + CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> "0.002", + CometConf.COMET_RESPECT_DATAFUSION_CONFIGS.key -> "true", + "spark.comet.datafusion.execution.sort_spill_reservation_bytes" -> "65536") ++ + codecConf: _*) { + val sorted = sql("SELECT * FROM tbl").sortWithinPartitions($"_1".desc) + assert(sorted.collect().length == numRows) + val plan = sorted.queryExecution.executedPlan + val sorts = collect(plan) { case sort: CometSortExec => sort } + assert(sorts.nonEmpty, s"Expected a native sort:\n$plan") + spilledBytes = sorts.map(_.metrics("spilled_bytes").value).sum + assert(spilledBytes > 0L, "Native sort did not spill") + } + spilledBytes + } + + // `spilled_bytes` counts what DataFusion wrote to disk, so this shows the default codec + // reaching DataFusion's spill writer, not just crossing JNI. + val uncompressed = + sortSpilledBytes(CometConf.COMET_EXEC_SPILL_COMPRESSION_CODEC.key -> "none") + val compressed = sortSpilledBytes() + assert( + compressed < uncompressed, + s"Spilled $compressed bytes with the default codec and $uncompressed with none") + } + } + } + test("sample without replacement") { withParquetTable((0 until 1000).map(i => (i, i + 1)), "tbl") { val df = sql("SELECT * FROM tbl").sample(withReplacement = false, fraction = 0.3, seed = 42) @@ -2854,6 +2900,14 @@ class CometExecSuite extends CometTestBase { } } + test("expand operator with no output columns falls back to Spark") { + // Column pruning leaves this shape when nothing above the Expand reads its output. + val child = spark.range(3).queryExecution.analyzed + val expand = Expand(Seq(Seq.empty[Expression], Seq.empty[Expression]), Seq.empty, child) + val df = datasetOfRows(spark, expand).groupBy().count() + checkSparkAnswerAndFallbackReason(df, "Expand without output columns is not supported") + } + test("multiple distinct multiple columns sets") { withTable("agg2") { val data2 = Seq[(Integer, Integer, Integer)]( @@ -3970,6 +4024,55 @@ class CometExecSuite extends CometTestBase { }) } + // https://github.com/apache/datafusion-comet/issues/6234. The final TopK reads its child's + // batches through an Arrow C stream, and Arrow Java hands native only the text of an exception + // thrown while producing one. The first batch is read on the JVM to derive the stream's schema, + // so the overflow has to land past it: row 90000 of a single file, read by one task. + test("TakeOrderedAndProjectExec: an input error past the first batch keeps Spark's exception") { + withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") { + withTempPath { dir => + spark + .range(0, 100000, 1, 1) + .selectExpr("CAST(id AS INT) AS id") + .write + .parquet(dir.getCanonicalPath) + spark.read.parquet(dir.getCanonicalPath).createOrReplaceTempView("overflow_src") + val df = + sql(s"SELECT id + ${Int.MaxValue - 90000} AS v FROM overflow_src ORDER BY v LIMIT 5") + checkSparkError(df, "ARITHMETIC_OVERFLOW") + } + } + } + + // https://github.com/apache/datafusion-comet/issues/6234. CometSparkToColumnarExec exports its + // rows to native without reading a batch on the JVM first, so even an exception thrown for the + // first batch reached the user only as the text inside a CometNativeException. + test("SparkToColumnar keeps the exception its row input throws") { + val rows = spark.sparkContext.parallelize(1 to 10, 1).map { i => + if (i == 5) throw new IllegalStateException("injected row input failure") + Row(i) + } + val df = spark + .createDataFrame(rows, StructType(Seq(StructField("a", IntegerType)))) + .groupBy() + .sum("a") + assert( + collect(df.queryExecution.executedPlan) { case c: CometSparkToColumnarExec => c }.nonEmpty, + "expected the rows to reach native through CometSparkToColumnarExec:\n" + + df.queryExecution.executedPlan) + + val (sparkError, cometError) = checkSparkAnswerMaybeThrows(df) + def injected(error: Option[Throwable], engine: String): (Class[_], Int) = { + val chain = causeChain(error.getOrElse(fail(s"$engine did not fail"))) + assert(!chain.exists(_.isInstanceOf[CometNativeException]), s"$engine: ${chain.head}") + val depth = chain.indexWhere(t => + t.isInstanceOf[IllegalStateException] && t.getMessage == "injected row input failure") + assert(depth >= 0, s"$engine did not surface the injected failure: ${chain.head}") + (chain.head.getClass, depth) + } + assert(injected(cometError, "Comet") == injected(sparkError, "Spark")) + } + // Arrow Java ignores ArrowArray.offset on import, and a sliced boolean is the one array arrow-rs // exports with a non-zero offset, so native has to zero it at every level, not only the top. // https://github.com/apache/datafusion-comet/issues/6288 @@ -4610,6 +4713,52 @@ class CometExecSuite extends CometTestBase { } } + // https://github.com/apache/datafusion-comet/issues/6573 + private val inputFileColumns = + Seq("input_file_name()", "input_file_block_start()", "input_file_block_length()", "id") + private val inputFileFallbackReason = + "Spark to Arrow conversion is not compatible with input_file_name" + + test("input_file_name above a converted Spark file scan keeps the scan on Spark") { + for (format <- Seq("parquet", "csv", "json"); v1List <- Seq("", format)) { + withTempPath { dir => + spark.range(9000).repartition(3).write.format(format).save(dir.toString) + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> v1List, + CometConf.COMET_CONVERT_FROM_PARQUET_ENABLED.key -> "true", + CometConf.COMET_CONVERT_FROM_CSV_ENABLED.key -> "true", + CometConf.COMET_CONVERT_FROM_JSON_ENABLED.key -> "true") { + // The filter would run in Comet above the conversion. Without a Comet operator there, + // Comet drops the conversion and Spark's project reads the scan directly. + val df = spark.read + .schema("id LONG") + .format(format) + .load(dir.toString) + .where("id >= 0") + .selectExpr(inputFileColumns: _*) + val (_, cometPlan) = checkSparkAnswerAndFallbackReason(df, inputFileFallbackReason) + assert(collect(cometPlan) { case c: CometSparkToColumnarExec => c }.isEmpty) + } + } + } + } + + test("input_file_name above a converted RDD scan keeps the scan on Spark") { + withTempPath { dir => + spark.range(9000).repartition(3).selectExpr("CAST(id AS STRING)").write.text(dir.toString) + withSQLConf(CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> "true") { + // HadoopRDD sets the file Spark evaluates input_file_name against, as FileScanRDD does. + val rows = spark.sparkContext.textFile(dir.toString).map(line => Row(line.toLong)) + val df = spark + .createDataFrame(rows, new StructType().add("id", LongType)) + .where("id >= 0") + .selectExpr(inputFileColumns: _*) + val (_, cometPlan) = checkSparkAnswerAndFallbackReason(df, inputFileFallbackReason) + assert(collect(cometPlan) { case c: CometSparkToColumnarExec => c }.isEmpty) + } + } + } + test("read CSV file") { Seq("", "csv").foreach { v1List => withSQLConf( diff --git a/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala index 1ccfacdffe5..7feefee3b26 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala @@ -224,19 +224,76 @@ class CometGenerateExecSuite extends CometTestBase { } } - test("explode with map input falls back") { - withSQLConf( - CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", - CometConf.COMET_EXEC_EXPLODE_ENABLED.key -> "true") { - val df = Seq((1, Map("a" -> 1, "b" -> 2)), (2, Map("c" -> 3))) - .toDF("id", "map") - .selectExpr("id", "explode(map) as (key, value)") - checkSparkAnswerAndFallbackReason( - df, - "Comet only supports explode/explode_outer for arrays, not maps") + for (generator <- Seq("explode", "explode_outer", "posexplode", "posexplode_outer")) { + test(s"$generator with JVM-fed map input") { + withSQLConf( + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", + CometConf.COMET_EXEC_EXPLODE_ENABLED.key -> "true") { + val input = Seq( + (1, Map("b" -> Integer.valueOf(20), "a" -> Integer.valueOf(10))), + (2, Map.empty[String, Integer]), + (3, null.asInstanceOf[Map[String, Integer]]), + (4, Map("null" -> null.asInstanceOf[Integer]))) + .toDF("id", "m") + checkMapGenerator(input, "m", generator) + } + } + + test(s"$generator with map input across batch boundaries") { + withSQLConf( + CometConf.COMET_EXEC_EXPLODE_ENABLED.key -> "true", + CometConf.COMET_BATCH_SIZE.key -> "4") { + val rows = (0 until 16).map { i => + val m = i % 4 match { + case 0 => null.asInstanceOf[Map[Int, java.lang.Integer]] + case 1 => Map.empty[Int, java.lang.Integer] + case _ => + (0 until 13).map { j => + j -> (if (j % 3 == 0) null else java.lang.Integer.valueOf(i * 100 + j)) + }.toMap + } + val booleans = Option(m) + .map(_.map { case (key, value) => + // A period of five cannot hide a wrong bitmap offset at four-row batch boundaries. + key -> (if (value == null) null else java.lang.Boolean.valueOf((i + key) % 5 < 2)) + }) + .orNull + (i, m, booleans) + } + withParquetDataFrame(rows) { input => + // One map exceeds the output batch size. Carry the map through too, so + // outer padding cannot accidentally replace the original empty map. + val maps = input.toDF("id", "ints", "booleans") + Seq("ints", "booleans").foreach { name => + withClue(s"$name: ") { checkMapGenerator(maps, name, generator) } + } + } + } } } + for (generator <- Seq("explode_outer", "posexplode_outer")) { + test(s"$generator with batches containing only null and empty maps") { + withSQLConf( + CometConf.COMET_EXEC_EXPLODE_ENABLED.key -> "true", + CometConf.COMET_BATCH_SIZE.key -> "4") { + val rows = (0 until 16).map { i => + val m = if (i % 2 == 0) null else Map.empty[Int, java.lang.Boolean] + (i, m) + } + withParquetDataFrame(rows) { input => + checkMapGenerator(input.toDF("id", "m"), "m", generator) + } + } + } + } + + private def checkMapGenerator(input: DataFrame, name: String, generator: String): Unit = { + val query = input.selectExpr("id", name, s"$generator($name)") + checkSparkAnswerAndOperator(query, Seq(classOf[CometExplodeExec])) + checkSparkSchema(query) + } + test("explode with nullable projected column") { withSQLConf( CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", @@ -390,16 +447,16 @@ class CometGenerateExecSuite extends CometTestBase { } } - test("posexplode with map input falls back") { + test("posexplode with map input falls back when disabled") { withSQLConf( CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", - CometConf.COMET_EXEC_EXPLODE_ENABLED.key -> "true") { + CometConf.COMET_EXEC_EXPLODE_ENABLED.key -> "false") { val df = Seq((1, Map("a" -> 1, "b" -> 2)), (2, Map("c" -> 3))) .toDF("id", "map") .selectExpr("id", "posexplode(map) as (pos, key, value)") checkSparkAnswerAndFallbackReason( df, - "Comet only supports explode/explode_outer for arrays, not maps") + "Native support for operator GenerateExec is disabled") } } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index e161466c4fb..aba97d80219 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -1437,6 +1437,12 @@ class CometInMemoryCacheSuite extends CometTestBase { extraConfs.containsKey(serializerKey) } + test("Comet's cache format is the default from Spark 3.5 and opt-in on Spark 3.4") { + // Spark 3.4 has no hook for CometCoalesceShufflePartitions, so AQE could not coalesce the + // shuffled branch of a Comet union whose other branch reads a cache in Comet's format. + assert(CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.defaultValue.get == isSpark35Plus) + } + test("Comet plugin installs its cache serializer only if Comet can scan the cache natively") { val cometOn = CometConf.COMET_ENABLED.key -> "true" val execOn = CometConf.COMET_EXEC_ENABLED.key -> "true" diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index eec16205a11..369c2517636 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -31,16 +31,20 @@ import org.apache.spark.sql.catalyst.TableIdentifier import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, DynamicPruningExpression, IsNotNull} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight} -import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortMergeJoinExec, CometUnionExec} -import org.apache.spark.sql.execution.{LocalTableScanExec, SparkPlan} -import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec} +import org.apache.spark.sql.catalyst.plans.Inner +import org.apache.spark.sql.catalyst.plans.logical.Join +import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometBroadcastHashJoinExec, CometBroadcastNestedLoopJoinExec, CometFilterExec, CometHashJoinExec, CometNativeScanExec, CometSortExec, CometSortMergeJoinExec, CometUnionExec, CometWindowExec} +import org.apache.spark.sql.execution.{ColumnarToRowTransition, InputAdapter, LocalTableScanExec, SortExec, SparkPlan, WholeStageCodegenExec} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, QueryStageExec} import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeLike} +import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, IntegerType, MetadataBuilder, StructField, StructType} -import org.apache.comet.CometConf -import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus +import org.apache.comet.{CometConf, CometExplainInfo, ExtendedExplainInfo} +import org.apache.comet.CometSparkSessionExtensions.{hasFallbackReason, isSpark35Plus} +import org.apache.comet.rules.RewriteJoin class CometJoinSuite extends CometTestBase { @@ -890,6 +894,331 @@ class CometJoinSuite extends CometTestBase { } } + // The MERGE hint makes Spark plan a SortMergeJoin, which forceShuffledHashJoin then rewrites. + private val forcedSmjQuery = + "SELECT /*+ MERGE(tbl_a) */ * FROM tbl_a JOIN tbl_b ON tbl_a._2 = tbl_b._1" + + // tbl_b is the smaller side, so the rewrite builds on the right. + private def withForcedSmjTables(f: => Unit): Unit = { + withParquetTable((0 until 1000).map(i => (i, i % 5)), "tbl_a") { + withParquetTable((0 until 10).map(i => (i % 10, i + 2)), "tbl_b") { + f + } + } + } + + private def rewriteReason(reason: String): String = + s"Cannot rewrite SortMergeJoin to HashJoin: $reason" + + // Checks that the SortMergeJoin and its sorts survived and that the kept join, which still + // runs natively, is not reported as a fallback. Returns Spark's size estimate of the build + // side, read through the same logical link the rewrite used. + private def assertSortMergeJoinKept(plan: SparkPlan): BigInt = { + val joins = collect(plan) { case j: CometSortMergeJoinExec => j } + assert(joins.size == 1, plan) + assert(collect(plan) { case s: CometSortExec => s }.size == 2, plan) + assert(collect(plan) { case j: CometHashJoinExec => j }.isEmpty, plan) + val kept = joins.head + assert(!hasFallbackReason(kept) && !hasFallbackReason(kept.originalPlan), plan) + val reasons = new ExtendedExplainInfo().getFallbackReasons(plan) + assert(!reasons.exists(_.contains("SortMergeJoin")), reasons) + kept.originalPlan.logicalLink match { + case Some(join: Join) => join.right.stats.sizeInBytes + case other => fail(s"Expected a Join logical link, got $other") + } + } + + // The refusal is recorded as information on the kept join and rendered as COMET-INFO. + private def assertKeptWithInfo(plan: SparkPlan, expected: String): Unit = { + val join = collect(plan) { case j: CometSortMergeJoinExec => j }.head + assert(join.getTagValue(CometExplainInfo.EXTENSION_INFO).exists(_.contains(expected)), plan) + val verbose = new ExtendedExplainInfo().generateVerboseInfo(plan) + assert(verbose.contains("[COMET-INFO:") && verbose.contains(expected), verbose) + } + + private def assertRewrittenToHashJoin(plan: SparkPlan): Unit = { + val joins = collect(plan) { case j: CometHashJoinExec => j } + assert(joins.map(_.buildSide) == Seq(BuildRight), plan) + assert(collect(plan) { case s: CometSortExec => s }.isEmpty, plan) + assert(collect(plan) { case j: CometSortMergeJoinExec => j }.isEmpty, plan) + } + + private def assertKeptWithReason(plan: SparkPlan, limit: Long): Unit = { + val estimate = assertSortMergeJoinKept(plan) + assert(estimate >= limit, estimate) + assertKeptWithInfo( + plan, + rewriteReason( + s"build side size estimate of $estimate bytes is not under the limit of $limit bytes")) + } + + for (adaptive <- Seq(false, true)) { + test(s"forceShuffledHashJoin keeps SortMergeJoin over maxBuildSize, AQE=$adaptive") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "1") { + withForcedSmjTables { + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) + assertKeptWithReason(cometPlan, 1) + } + } + } + + test(s"forceShuffledHashJoin rewrites SortMergeJoin under maxBuildSize, AQE=$adaptive") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "1g") { + withForcedSmjTables { + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) + assertRewrittenToHashJoin(cometPlan) + } + } + } + } + + test("forceShuffledHashJoin matches Spark's strict size rule at the limit") { + // AQE off, so the estimate read from the first plan is the one the later runs compare. + withForcedSmjTables { + var estimate = BigInt(0) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "1") { + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) + estimate = assertSortMergeJoinKept(cometPlan) + } + // A build side exactly at the limit is not under it, as in canBuildLocalHashMapBySize. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> estimate.toString) { + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) + assertKeptWithReason(cometPlan, estimate.toLong) + } + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> (estimate + 1).toString) { + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) + assertRewrittenToHashJoin(cometPlan) + } + } + } + + test("forceShuffledHashJoin with a non-positive maxBuildSize rewrites regardless of size") { + withSQLConf( + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1") { + withForcedSmjTables { + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) + assertRewrittenToHashJoin(cometPlan) + } + } + } + + test("forceShuffledHashJoin without maxBuildSize uses Spark's shuffled hash join size rule") { + // The limit is autoBroadcastJoinThreshold times the shuffle partition count. + withSQLConf( + CometConf.COMET_FORCE_SHJ.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1") { + withForcedSmjTables { + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) + assertKeptWithReason(cometPlan, 2) + } + } + withSQLConf( + CometConf.COMET_FORCE_SHJ.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1g") { + withForcedSmjTables { + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) + assertRewrittenToHashJoin(cometPlan) + } + } + } + + test("forceShuffledHashJoin uses Spark's default threshold when broadcasts are disabled") { + // Disabling broadcasts says nothing about the hash table an executor can hold, so the limit + // falls back to Spark's default threshold times the partition count. AQE is off so the + // decision uses the planning estimate, which fileCompressionFactor scales past 10 MB without + // writing a 10 MB table. + val defaultThreshold = SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.defaultValue.get + assert(defaultThreshold == 10L * 1024 * 1024) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_FORCE_SHJ.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + withForcedSmjTables { + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) + assertRewrittenToHashJoin(cometPlan) + } + withSQLConf(SQLConf.FILE_COMPRESSION_FACTOR.key -> "1000000") { + withForcedSmjTables { + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(forcedSmjQuery)) + assertKeptWithReason(cometPlan, defaultThreshold) + } + } + } + } + + test("forceShuffledHashJoin keeps a SortMergeJoin that has no statistics") { + val left = spark.range(1).toDF("a").queryExecution.sparkPlan + val right = spark.range(1).toDF("b").queryExecution.sparkPlan + def smj: SortMergeJoinExec = + SortMergeJoinExec(left.output, right.output, Inner, None, left, right) + assert(smj.logicalLink.isEmpty) + + val kept = smj + assert(RewriteJoin.rewrite(kept, spark.sessionState.conf) eq kept) + assert(!hasFallbackReason(kept)) + val info = kept.getTagValue(CometExplainInfo.EXTENSION_INFO) + assert( + info.exists(_.contains(rewriteReason("no statistics are available for the build side"))), + info) + + withSQLConf(CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1") { + val rewritten = RewriteJoin.rewrite(smj, spark.sessionState.conf) + assert(rewritten.isInstanceOf[ShuffledHashJoinExec], rewritten) + } + } + + // Both joins are on the same key, so EnsureRequirements puts no sort between them: the inner + // sort-merge join's output ordering satisfies the outer one. When the inner join is rewritten + // to a hash join and the outer one is kept, the outer join needs that ordering restored. + private def withChainedJoinTables(midRows: Int, midKeys: Int)(f: => Unit): Unit = { + withParquetTable((0 until 10000).map(i => (i % 100, i)), "big") { + withParquetTable((0 until 10).map(i => (i * 10, i)), "small") { + withParquetTable((0 until midRows).map(i => (i % midKeys, i)), "mid") { + f + } + } + } + } + + // Whether a hash join is reachable without leaving the stage. + private def reachesHashJoinInStage(plan: SparkPlan): Boolean = plan match { + case _: CometHashJoinExec => true + case _: ShuffleExchangeLike | _: QueryStageExec => false + case p => p.children.exists(reachesHashJoinInStage) + } + + // Whether a sort-merge join input is a sort over a hash join in the same stage. A Spark + // sort-merge join reads the sort through codegen wrappers and a columnar to row transition. + private def isSortOverHashJoin(input: SparkPlan): Boolean = input match { + case s: CometSortExec => reachesHashJoinInStage(s.child) + case s: SortExec => reachesHashJoinInStage(s.child) + case t: ColumnarToRowTransition => isSortOverHashJoin(t.child) + case w: WholeStageCodegenExec => isSortOverHashJoin(w.child) + case a: InputAdapter => isSortOverHashJoin(a.child) + case _ => false + } + + // Checks the answer first, then that a kept sort-merge join reads the rewritten hash join + // through a sort. + private def checkSortRestoredOverHashJoin(query: String): Unit = { + val (_, cometPlan) = checkSparkAnswer(query) + val inputs = collect(cometPlan) { + case j: CometSortMergeJoinExec => j.children + case j: SortMergeJoinExec => j.children + }.flatten + assert(inputs.exists(isSortOverHashJoin), cometPlan) + } + + private val chainedJoinQuery = "SELECT big._1, big._2, small._2, mid._2 FROM big " + + "JOIN small ON big._1 = small._1 JOIN mid ON big._1 = mid._1" + + test("forceShuffledHashJoin restores ordering for a kept SortMergeJoin, AQE=false") { + // The default limit, 3000 bytes times two partitions, falls between the two build sides. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_FORCE_SHJ.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "3000") { + withChainedJoinTables(midRows = 3000, midKeys = 100) { + checkSortRestoredOverHashJoin(chainedJoinQuery) + } + } + } + + test("forceShuffledHashJoin restores ordering for a kept SortMergeJoin, AQE=true") { + // Under AQE the decision uses the materialized shuffle sizes, so mid is made large enough + // to stay over the limit while small stays under it. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "1000000", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "3000") { + withChainedJoinTables(midRows = 300000, midKeys = 30000) { + checkSortRestoredOverHashJoin(chainedJoinQuery) + } + } + } + + // The rewrite never converts a LeftSemi or ExistenceJoin that builds on the right, whatever + // the build size, so these joins are kept above the rewritten inner join with no size limit. + // Each predicate also refers to small, so neither the filter nor the semi join is pushed below + // the inner join. + for (adaptive <- Seq(false, true)) { + def withKeptSemiJoinConf(f: => Unit): Unit = withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1") { + withChainedJoinTables(midRows = 3000, midKeys = 100)(f) + } + + test(s"forceShuffledHashJoin restores ordering for a kept ExistenceJoin, AQE=$adaptive") { + withKeptSemiJoinConf { + checkSortRestoredOverHashJoin( + "SELECT big._1, big._2, small._2 FROM big JOIN small ON big._1 = small._1 " + + "WHERE small._2 < 0 OR EXISTS (SELECT 1 FROM mid WHERE mid._1 = big._1)") + } + } + + test(s"forceShuffledHashJoin restores ordering for a kept LeftSemi join, AQE=$adaptive") { + withSQLConf( + SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> + "org.apache.spark.sql.catalyst.optimizer.PushLeftSemiLeftAntiThroughJoin") { + withKeptSemiJoinConf { + checkSortRestoredOverHashJoin( + "SELECT big._1, big._2, small._2 FROM big JOIN small ON big._1 = small._1 " + + "WHERE EXISTS (SELECT 1 FROM mid WHERE mid._1 = big._1 AND mid._2 > small._2)") + } + } + } + } + + // A window partitioned by the join key gets no sort of its own, since the sort-merge join's + // output ordering satisfies it, so the window needs that ordering restored too. + for (adaptive <- Seq(false, true)) { + test(s"forceShuffledHashJoin restores ordering for a window over the join, AQE=$adaptive") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + CometConf.COMET_FORCE_SHJ.key -> "true", + CometConf.COMET_FORCE_SHJ_MAX_BUILD_SIZE.key -> "-1") { + withChainedJoinTables(midRows = 3000, midKeys = 100) { + val (_, cometPlan) = checkSparkAnswer( + "SELECT big._1, big._2, small._2, count(*) OVER (PARTITION BY big._1) " + + "FROM big JOIN small ON big._1 = small._1") + assert(collect(cometPlan) { case j: CometHashJoinExec => j }.nonEmpty, cometPlan) + assert(collect(cometPlan) { case w: CometWindowExec => w }.nonEmpty, cometPlan) + } + } + } + } + test("BroadcastHashJoin with LeftAnti and NOT IN subquery (null-aware)") { withSQLConf( SQLConf.PREFER_SORTMERGEJOIN.key -> "false", diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeReaderSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeReaderSuite.scala index 63a7c3f0420..d5dff80cb2f 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeReaderSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeReaderSuite.scala @@ -1581,6 +1581,92 @@ class CometNativeReaderSuite extends CometTestBase with AdaptiveSparkPlanHelper } } + test("row-group statistics pruning fires for a floating-point comparison") { + // Native comparisons normalize float operands to follow Spark's ordering, which pruning + // cannot see through. A scan's data filters are built without that normalization, so they + // still prune, and the Filter above the scan applies Spark's semantics to the rows. Spark + // orders NaN above every other value, so `d > 500.0D` matches the NaN in the first row + // group even though every other value there is smaller: pruning must keep that row group. + withTempPath { dir => + withSQLConf(SQLConf.LEAF_NODE_DEFAULT_PARALLELISM.key -> "1") { + spark + .range(0, 1000) + .selectExpr("IF(id = 10, double('NaN'), CAST(id AS DOUBLE)) AS d") + .repartition(1) + .write + .option("parquet.block.size", "1024") + .format("parquet") + .save(dir.toString) + + val parquetFile = dir + .listFiles() + .find(_.getName.endsWith(".parquet")) + .getOrElse(fail("No parquet file was written")) + val reader = ParquetFileReader.open( + org.apache.parquet.hadoop.util.HadoopInputFile + .fromPath(new Path(parquetFile.getAbsolutePath), spark.sessionState.newHadoopConf())) + val numRowGroups = + try reader.getRowGroups.size() + finally reader.close() + assert(numRowGroups > 1, s"Test setup needs >1 row groups, got $numRowGroups") + + val df = spark.read.parquet(dir.toString).where("d > 500.0D") + val (_, cometPlan) = checkSparkAnswerAndOperator(df) + val nativeScans = cometPlan.collect { case n: CometNativeScanExec => n } + assert(nativeScans.nonEmpty, "Expected a CometNativeScanExec") + val metrics = nativeScans.head.metrics + val pruned = metrics("row_groups_pruned_statistics").value + val matched = metrics("row_groups_matched_statistics").value + assert( + pruned > 0 && pruned + matched == numRowGroups, + "Row-group statistics pruning did not fire " + + s"(pruned=$pruned, matched=$matched of $numRowGroups total)") + assert(df.collect().exists(_.getDouble(0).isNaN), "The NaN row was filtered out") + } + } + } + + test("bloom filter pruning matches either signed zero for a floating-point equality") { + // A bloom filter holds the bits of each value, so `-0.0` and `0.0` are separate entries, + // while Spark's `=` matches either zero. The scan's data filter probes for both, so the file + // holding `-0.0` and the file holding `0.0` are both read for either zero. Spark's own reader + // probes with the literal's bits and skips the file holding the other zero, so the expected + // rows are written out instead of compared with Spark's. + import testImplicits._ + withTempPath { dir => + val path = dir.getCanonicalPath + Seq(Seq(-1.0, -0.0, 1.0), Seq(-2.0, 0.0, 2.0)).foreach { values => + values + .toDF("d") + .coalesce(1) + .write + .mode("append") + .option("parquet.bloom.filter.enabled#d", "true") + .option("parquet.enable.dictionary", "false") + .parquet(path) + } + + def bloomFilterPruning(df: DataFrame): (Long, Long) = { + val scans = collect(df.queryExecution.executedPlan) { case s: CometNativeScanExec => s } + assert(scans.size == 1, s"Expected one CometNativeScanExec:\n${df.queryExecution}") + val metrics = scans.head.metrics + ( + metrics("row_groups_pruned_bloom_filter").value, + metrics("row_groups_matched_bloom_filter").value) + } + + Seq("-0.0D", "0.0D").foreach { zero => + val df = spark.read.parquet(path).where(s"d = $zero") + checkAnswer(df, Seq(Row(-0.0), Row(0.0))) + assert(bloomFilterPruning(df) == ((0L, 2L)), s"d = $zero") + } + // The statistics of both files cover 0.5, so only their bloom filters prune them. + val df = spark.read.parquet(path).where("d = 0.5D") + checkAnswer(df, Seq.empty[Row]) + assert(bloomFilterPruning(df) == ((2L, 0L)), "d = 0.5D") + } + } + test("datafusion escape hatch pruning=false disables row-group statistics pruning") { // Regression test pinning the `pruning` field of the newly-plumbed session // `ParquetOptions` at the level users care about: an explicit diff --git a/spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala new file mode 100644 index 00000000000..03181136005 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala @@ -0,0 +1,428 @@ +/* + * 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.sql.{Date, Timestamp} +import java.time.LocalDateTime + +import org.apache.spark.sql.{CometTestBase, DataFrame, Row} +import org.apache.spark.sql.comet.{CometHashAggregateExec, CometPlan, CometSparkToColumnarExec} +import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} +import org.apache.spark.sql.execution.{ColumnarToRowExec, ColumnarToRowTransition, SparkPlan} +import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, SortAggregateExec} +import org.apache.spark.sql.functions.{array, avg, col, count, hash, length, max, min, size, spark_partition_id, sum} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, CalendarIntervalType, DataType, DataTypes, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, StructType, TimestampNTZType, TimestampType} +import org.apache.spark.unsafe.types.CalendarInterval + +import org.apache.comet.{CometConf, ExtendedExplainInfo} + +// Top-level, so the encoder needs no outer pointer. +case class ShuffleInputRec(a: Int, b: String) + +/** Tests for [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. */ +class CometShuffleInputConversionSuite extends CometTestBase { + + import testImplicits._ + + /** + * `CometTestBase` turns on the conversion of leaf operators such as RDD scans, which would + * convert the scans below the shuffles here before this conversion sees them. It is off by + * default. + */ + private def withoutLeafConversion(f: => Unit): Unit = + withSQLConf(sparkToArrowConversionConfs(enabled = false): _*)(f) + + /** Defines a test that runs with the shuffle input conversion enabled. */ + private def convertTest(name: String)(f: => Unit): Unit = + test(name) { + withoutLeafConversion { + withSQLConf(CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.key -> "true")(f) + } + } + + private val rowSchema = new StructType() + .add("k", IntegerType) + .add("l", LongType) + .add("s", DataTypes.StringType) + .add("d", DataTypes.DoubleType) + .add("m", DecimalType(18, 2)) + + /** + * `rows` rows from an RDD, so the shuffle above them reads a Spark operator: the conversion of + * RDD scans, `spark.comet.convert.rdd.enabled`, is off by default. + */ + private def rowsDf(rows: Int = 1000): DataFrame = { + val data = (0 until rows).map { i => + Row( + i % 23, + i.toLong, + if (i % 11 == 0) null else s"value-$i", + i + 0.5d, + java.math.BigDecimal.valueOf(i * 7919L % 1000000L, 2)) + } + spark.createDataFrame(spark.sparkContext.parallelize(data, 4), rowSchema) + } + + /** Rows with an `array` column, which `CometSparkToColumnarExec` does not convert. */ + private def arraysDf(rows: Int = 100): DataFrame = { + val schema = new StructType().add("k", IntegerType).add("xs", ArrayType(IntegerType)) + val data = (0 until rows).map(i => Row(i % 23, Seq(i, i + 1))) + spark.createDataFrame(spark.sparkContext.parallelize(data, 4), schema) + } + + private def conversions(plan: SparkPlan): Seq[CometSparkToColumnarExec] = + collectWithSubqueries(plan) { case c: CometSparkToColumnarExec => c } + + private def cometShuffles(plan: SparkPlan): Seq[CometShuffleExchangeExec] = + collectWithSubqueries(plan) { case s: CometShuffleExchangeExec => s } + + /** The native shuffles that read rows `CometSparkToColumnarExec` converted. */ + private def convertedShuffles(plan: SparkPlan): Seq[CometShuffleExchangeExec] = + cometShuffles(plan).filter { s => + s.shuffleType == CometNativeShuffle && s.child.isInstanceOf[CometSparkToColumnarExec] + } + + /** + * Spark inserts no columnar transitions below a `CometSparkToColumnarExec`, so the rule adds + * them for the operators under the conversion. Without one, an operator reads its Comet child + * through `CometExec.doExecute`, Spark's interpreted columnar-to-row path, which gives the + * right answer slowly, so only the plan shows it. + */ + private def assertRowOperatorsReadThroughTransitions(plan: SparkPlan): Unit = { + val converted = conversions(plan) + assert(converted.nonEmpty, plan) + converted.foreach { conversion => + // `InputAdapter` and `WholeStageCodegenExec` report their child's `supportsColumnar`. + val bare = conversion.child.collect { + case p + if !p.supportsColumnar && !p.isInstanceOf[ColumnarToRowTransition] && + p.children.exists(_.supportsColumnar) => + p + } + assert(bare.isEmpty, s"row operators read a columnar child without a transition:\n$plan") + assert( + conversion.child.collect { case c: ColumnarToRowExec => c }.isEmpty, + s"expected Comet's columnar-to-row transitions, not Spark's:\n$plan") + } + // The transitions are added on the first of the rule's passes under AQE. A later pass must + // not report them as operators Comet failed to convert. + val reasons = new ExtendedExplainInfo().getFallbackReasons(plan) + assert(!reasons.exists(_.contains("ColumnarToRow")), reasons) + } + + /** Runs `f` with each join planned as a sort-merge join of two shuffles into 10 partitions. */ + private def withShuffledJoins(f: => Unit): Unit = + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "10")(f) + + /** + * Values of each type the conversion admits as a hash key, with boundary values. The floating + * point values include NaN and both zeros, which a join normalizes. + */ + private val hashKeys: Seq[(DataType, Seq[Any])] = Seq( + BooleanType -> Seq(true, false), + ByteType -> Seq(0.toByte, -1.toByte, Byte.MinValue, Byte.MaxValue), + ShortType -> Seq(0.toShort, -1.toShort, Short.MinValue, Short.MaxValue), + IntegerType -> Seq(0, -1, Int.MinValue, Int.MaxValue), + LongType -> Seq(0L, -1L, Long.MinValue, Long.MaxValue), + FloatType -> Seq( + 0.0f, + -0.0f, + 1.5f, + Float.NaN, + Float.MinPositiveValue, + Float.MaxValue, + Float.NegativeInfinity), + DoubleType -> Seq( + 0.0d, + -0.0d, + 1.5d, + Double.NaN, + Double.MinPositiveValue, + Double.MaxValue, + Double.NegativeInfinity), + DecimalType(18, 2) -> Seq("0.00", "-0.01", "9999999999999999.99", "-9999999999999999.99") + .map(new java.math.BigDecimal(_)), + DateType -> Seq("1970-01-01", "1969-12-31", "0001-01-01", "9999-12-31").map(Date.valueOf), + TimestampType -> Seq( + new Timestamp(0L), + new Timestamp(-1L), + Timestamp.valueOf("0001-01-01 00:00:00"), + Timestamp.valueOf("9999-12-31 23:59:59.999999")), + TimestampNTZType -> Seq( + LocalDateTime.of(1970, 1, 1, 0, 0), + LocalDateTime.of(1969, 12, 31, 23, 59, 59, 999999000), + LocalDateTime.of(1, 1, 1, 0, 0), + LocalDateTime.of(9999, 12, 31, 23, 59, 59, 999999000)), + BinaryType -> Seq(Array.emptyByteArray, Array[Byte](0), Array[Byte](-1, -128, 127))) + + /** Four rows of each value and of a NULL, as column `k`, over four input partitions. */ + private def keysDf(keyType: DataType, values: Seq[Any]): DataFrame = { + val keys = values :+ null + val data = (0 until keys.length * 4).map(i => Row(keys(i % keys.length), i.toLong)) + val schema = new StructType().add("k", keyType).add("v", LongType) + spark.createDataFrame(spark.sparkContext.parallelize(data, 4), schema) + } + + convertTest("the shuffle of a Spark operator's rows runs as native shuffle") { + Seq("true", "false").foreach { aqe => + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) { + val (_, plan) = checkSparkAnswer( + rowsDf() + .repartition(10, col("k")) + .groupBy("k") + .agg(sum("l"), sum(length(col("s"))), sum("d"), sum("m"))) + assert(convertedShuffles(plan).length == 1, s"AQE $aqe:\n$plan") + assert( + !cometShuffles(plan).exists(_.shuffleType == CometColumnarShuffle), + s"AQE $aqe:\n$plan") + } + } + } + + convertTest("Spark operators over Comet operators read them through transitions") { + withParquetTable((0 until 200).map(i => (i, (i % 13).toString)), "tbl") { + Seq("true", "false").foreach { aqe => + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) { + val ds = spark.sql("SELECT _1 AS a, _2 AS b FROM tbl").as[ShuffleInputRec] + // The shuffle reads the typed operation's SerializeFromObject. + val (_, repartitioned) = checkSparkAnswer( + ds.map(r => ShuffleInputRec(r.a % 13, r.b)) + .repartition(7, col("a")) + .groupBy("a") + .agg(count("b"))) + assert(convertedShuffles(repartitioned).length == 1, s"AQE $aqe:\n$repartitioned") + assertRowOperatorsReadThroughTransitions(repartitioned) + // The shuffle reads a Spark partial aggregate over the typed operation. The final + // aggregate stays on Spark too, so the shuffle would go back to Spark's own. + withSQLConf(CometConf.COMET_SHUFFLE_REVERT_REDUNDANT_COLUMNAR_ENABLED.key -> "false") { + val (_, aggregated) = + checkSparkAnswer(ds.map(r => ShuffleInputRec(r.a % 7, r.b)).groupBy("a").count()) + assert(convertedShuffles(aggregated).length == 1, s"AQE $aqe:\n$aggregated") + assertRowOperatorsReadThroughTransitions(aggregated) + } + } + } + } + } + + convertTest("aggregates whose partial aggregate runs on Spark") { + // The partial aggregate reads Spark rows, so it runs on Spark, and native shuffle carries its + // buffers to the final aggregate. With these aggregates the final one runs natively. + val (_, split) = checkSparkAnswer( + rowsDf() + .groupBy((col("k") % 7).as("g")) + .agg(sum("l"), avg("d"), max("l"), min("l"))) + assert(convertedShuffles(split).length == 1, split) + assert(collect(split) { case a: HashAggregateExec => a }.length == 1, split) + assert(collect(split) { case a: CometHashAggregateExec => a }.length == 1, split) + // The min of a string makes Spark plan sort aggregates, which stay on Spark, so the final + // aggregate reads the buffers back as rows. The shuffle stays native: the final aggregate + // reads it through a sort, and the revert in the next test only takes a shuffle that an + // aggregate reads directly. + val (_, sorted) = checkSparkAnswer( + rowsDf() + .groupBy((col("k") % 7).as("g")) + .agg(sum("l"), count("s"), avg("d"), sum("m"), max("l"), min("s"))) + assert(convertedShuffles(sorted).length == 1, sorted) + assert(collect(sorted) { case a: SortAggregateExec => a }.length == 2, sorted) + } + + convertTest("a shuffle between two Spark aggregates goes back to Spark's shuffle") { + withSQLConf(CometConf.COMET_EXEC_AGGREGATE_ENABLED.key -> "false") { + val (_, plan) = checkSparkAnswer(rowsDf().groupBy("k").agg(sum("l"))) + assert(cometShuffles(plan).isEmpty, plan) + withSQLConf(CometConf.COMET_SHUFFLE_REVERT_REDUNDANT_COLUMNAR_ENABLED.key -> "false") { + val (_, kept) = checkSparkAnswer(rowsDf().groupBy("k").agg(sum("l"))) + assert(convertedShuffles(kept).length == 1, kept) + } + } + } + + convertTest("a stage that the transition revert puts back on Spark keeps the JVM shuffle") { + // The typed operation reads the Comet scan through a transition. With no transitions allowed, + // the revert puts the stage below the shuffle back on Spark, the conversion included, and the + // shuffle goes back to the JVM columnar shuffle, which reads the stage's rows. + withParquetTable((0 until 200).map(i => (i, (i % 13).toString)), "tbl") { + Seq("true", "false").foreach { aqe => + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe, + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") { + val ds = spark.sql("SELECT _1 AS a, _2 AS b FROM tbl").as[ShuffleInputRec] + val (_, plan) = checkSparkAnswer( + ds.map(r => ShuffleInputRec(r.a % 13, r.b)) + .repartition(7, col("a")) + .groupBy("a") + .agg(max("b"))) + assert(conversions(plan).isEmpty, s"AQE $aqe:\n$plan") + val reverted = cometShuffles(plan).filter(_.shuffleType == CometColumnarShuffle) + assert(reverted.length == 1, s"AQE $aqe:\n$plan") + assert(!reverted.head.child.exists(_.isInstanceOf[CometPlan]), s"AQE $aqe:\n$plan") + } + } + } + } + + hashKeys.foreach { case (keyType, values) => + convertTest(s"native shuffle puts ${keyType.simpleString} keys where Spark does") { + withShuffledJoins { + val keys = keysDf(keyType, values) + // The partition of each row, against Spark's partitioner, NULL keys included. + val (_, placed) = checkSparkAnswer( + keys.repartition(10, col("k")).select(col("k"), col("v"), spark_partition_id())) + assert(convertedShuffles(placed).length == 1, placed) + // The left input converts and the right one, with its array column, does not, so + // matching keys meet only if both shuffles put them in the same partition. + val left = keys.select(col("k").as("lk"), col("v")) + val right = keys.select(col("k").as("rk"), array(col("v")).as("xs")) + val (_, joined) = checkSparkAnswer(left.join(right, col("lk") === col("rk"))) + assert(convertedShuffles(joined).length == 1, joined) + assert(cometShuffles(joined).count(_.shuffleType == CometColumnarShuffle) == 1, joined) + } + } + } + + convertTest("a shuffle that hashes a wide decimal stays on the JVM columnar shuffle") { + // Native shuffle hashes a decimal wider than 18 digits differently from Spark's partitioner + // (#5994). The right input, with its array column, stays on the JVM columnar shuffle, + // so the left one has to as well, or matching keys land in different partitions. + withShuffledJoins { + val left = rowsDf(100).select(col("k").cast(DecimalType(38, 0)).as("lk"), col("l")) + val right = arraysDf().select(col("k").cast(DecimalType(38, 0)).as("rk"), col("xs")) + val df = left.join(right, col("lk") === col("rk")) + val (_, plan) = checkSparkAnswer(df) + assert(conversions(plan).isEmpty, plan) + checkCometExchange(df, 2, native = false) + } + } + + convertTest("a shuffle that hashes a string stays on the JVM columnar shuffle") { + // Spark's partitioner hashes a string's bytes as they are, but native shuffle hashes them + // after the import into native has replaced invalid UTF-8. The right input, with its + // array column, stays on the JVM columnar shuffle, so the left one has to as well, or + // matching keys with invalid UTF-8 land in different partitions. A key computed from a + // string, such as its hash, is no different: native shuffle would compute it from the + // replaced string. Both shuffles write the rows with invalid UTF-8 replaced, and the join + // compares what they wrote, so the two keys stay apart only because each key expression puts + // them in different partitions. + withShuffledJoins { + val schema = new StructType().add("b", BinaryType).add("v", LongType) + val data = (0 until 40).map { i => + Row(Array((if (i % 2 == 0) 0x80 else 0x83).toByte), i.toLong) + } + val keys = spark.createDataFrame(spark.sparkContext.parallelize(data, 4), schema) + val left = keys.select(col("b").cast(StringType).as("lk"), col("v")) + val right = keys.select( + col("b").cast(StringType).as("rk"), + array(col("v").cast(IntegerType)).as("xs")) + Seq(col("lk") === col("rk"), hash(col("lk")) === hash(col("rk"))).foreach { condition => + val df = left.join(right, condition) + val (_, plan) = checkSparkAnswer(df) + assert(conversions(plan).isEmpty, s"$condition:\n$plan") + checkCometExchange(df, 2, native = false) + } + } + } + + convertTest("nested columns over several batches") { + // The rows of each input partition fill several batches, and the conversion reuses its + // vectors from one batch to the next. With the conversion of RDD scans on, native shuffle + // reads the converted scan the same way. + val schema = new StructType() + .add("k", IntegerType) + .add("payload", new StructType().add("v", LongType).add("s", StringType)) + .add("xs", ArrayType(StringType)) + .add("m", MapType(StringType, StringType)) + val data = (0 until 24).map { i => + Row( + i, + if (i % 5 == 4) null else Row(i.toLong, if (i % 3 == 0) null else s"s$i"), + if (i % 6 == 5) null else Seq.tabulate(i % 4)(j => if (j == 1) null else s"x$i-$j"), + if (i % 7 == 6) null + else Seq.tabulate(i % 3)(j => s"k$j" -> (if (j == 1) null else s"v$i-$j")).toMap) + } + Seq("false", "true").foreach { leafConversion => + withSQLConf( + CometConf.COMET_BATCH_SIZE.key -> "7", + CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> leafConversion) { + val df = spark.createDataFrame(spark.sparkContext.parallelize(data, 1), schema) + Seq(df.repartition(3, col("k")), df.repartitionByRange(3, col("k"))).foreach { shuffled => + val (_, plan) = checkSparkAnswer(shuffled) + assert(convertedShuffles(plan).length == 1, s"leaf conversion $leafConversion:\n$plan") + } + } + } + } + + convertTest("calendar intervals keep Spark's shuffle") { + // Arrow holds an interval's time part in nanoseconds, which can't hold every number of + // microseconds Spark can. The JVM columnar shuffle doesn't take calendar intervals either. + val schema = new StructType().add("k", IntegerType).add("i", CalendarIntervalType) + val data = Seq(Row(1, new CalendarInterval(0, 0, 10800000000000000L))) + val df = spark + .createDataFrame(spark.sparkContext.parallelize(data, 1), schema) + .repartition(2, col("k")) + val (_, plan) = checkSparkAnswer(df) + assert(cometShuffles(plan).isEmpty, plan) + } + + convertTest("columns the conversion does not support keep the JVM columnar shuffle") { + val df = arraysDf().repartition(5, col("k")).groupBy("k").agg(sum(size(col("xs")))) + val (_, plan) = checkSparkAnswer(df) + assert(conversions(plan).isEmpty, plan) + assert(cometShuffles(plan).forall(_.shuffleType == CometColumnarShuffle), plan) + } + + convertTest("range partitioning") { + val (_, plan) = checkSparkAnswer(rowsDf().orderBy(col("l").desc)) + assert(convertedShuffles(plan).length == 1, plan) + // The range partitioner samples the rows that the conversion reads, without converting them, + // so the conversion only counts the rows that the shuffle writes. + assert(conversions(plan).head.metrics("numInputRows").value == 1000, plan) + } + + convertTest("the conversion only takes over from the JVM columnar shuffle") { + // In `jvm` mode the shuffle keeps the JVM columnar shuffle. In `native` mode, where Comet has + // no shuffle for a Spark operator's rows, it keeps Spark's shuffle. + Seq("jvm" -> Seq(CometColumnarShuffle), "native" -> Seq.empty).foreach { + case (mode, shuffleTypes) => + withSQLConf(CometConf.COMET_SHUFFLE_MODE.key -> mode) { + val (_, plan) = + checkSparkAnswer(rowsDf().repartition(10, col("k")).groupBy("k").count()) + assert(conversions(plan).isEmpty, s"$mode:\n$plan") + assert( + cometShuffles(plan).map(_.shuffleType).distinct == shuffleTypes, + s"$mode:\n$plan") + } + } + } + + test("off by default") { + withoutLeafConversion { + val (_, plan) = checkSparkAnswer(rowsDf().repartition(10, col("k")).groupBy("k").count()) + assert(conversions(plan).isEmpty, plan) + assert(cometShuffles(plan).exists(_.shuffleType == CometColumnarShuffle), plan) + } + } +} diff --git a/spark/src/test/scala/org/apache/comet/exec/CometTypedDatasetSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometTypedDatasetSuite.scala index 9b809e24e8d..c9e1a349e3d 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometTypedDatasetSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometTypedDatasetSuite.scala @@ -340,6 +340,26 @@ class CometTypedDatasetSuite extends CometTestBase { } } + // https://github.com/apache/datafusion-comet/issues/6573 + convertTest("input_file_name above a map keeps the map's output on Spark") { + withTempPath { dir => + spark.range(9000).repartition(3).write.parquet(dir.toString) + // The filter would run in Comet above the conversion, which would read ahead of the files + // Spark's reader is on when Spark evaluates input_file_name in the project. + val df = spark.read + .parquet(dir.toString) + .as[Long] + .map(_ + 1) + .toDF("id") + .where("id >= 0") + .selectExpr("input_file_name()", "input_file_block_start()", "input_file_block_length()") + val (_, plan) = checkSparkAnswerAndFallbackReason( + df, + "Spark to Arrow conversion is not compatible with input_file_name") + assert(conversions(plan).isEmpty, plan) + } + } + test("off by default") { withRecs() { ds => val (_, plan) = checkSparkAnswer(ds.map(r => TypedDsRec(r.a + 1, r.b)).groupBy("b").count()) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala index 700748b6edc..51a918606ce 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala @@ -214,10 +214,10 @@ class CometWindowExecSuite extends CometTestBase { assume(isSpark35Plus, "WindowGroupLimit was added in Spark 3.5") // The native RANK and DENSE_RANK find ties by byte equality, on order keys that the native // planner has normalized, nested floats included. So [-0.0] and [0.0] are peers, and the - // cutoff keeps both rows, as Spark does (#5507). Strict floating-point mode declines a nested - // key whose type can hold a null element or field (#6476, #6477), and these columns, read - // back from Parquet, can. ROW_NUMBER never compares peers, and Spark normalizes nested - // floating-point partition keys itself. + // cutoff keeps both rows, as Spark does (#5507). These keys, read back from Parquet, can hold + // a null element, but their default null order places it where Spark does (#6476), so they + // stay native in strict floating-point mode too. ROW_NUMBER never compares peers, and Spark + // normalizes nested floating-point partition keys itself. withTempDir { dir => val path = new Path(dir.toString, "nested_float_order").toString Seq((1, -0.0f, -0.0d), (2, 0.0f, 0.0d), (3, 1.0f, 1.0d)) @@ -248,7 +248,7 @@ class CometWindowExecSuite extends CometTestBase { assert(rankLimit().collect().map(_.getInt(0)).sorted.toSeq == Seq(1, 2)) withSQLConf(CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "true") { - checkSparkAnswerAndFallbackReason(rankLimit(), "can hold a null element or field") + checkSparkAnswerAndOperator(rankLimit(), Seq(classOf[CometWindowGroupLimitExec])) assert(rankLimit().collect().map(_.getInt(0)).sorted.toSeq == Seq(1, 2)) } } diff --git a/spark/src/test/scala/org/apache/comet/iceberg/IcebergReflectionSuite.scala b/spark/src/test/scala/org/apache/comet/iceberg/IcebergReflectionSuite.scala index e7bff634db4..f5f3a75f493 100644 --- a/spark/src/test/scala/org/apache/comet/iceberg/IcebergReflectionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/iceberg/IcebergReflectionSuite.scala @@ -138,6 +138,27 @@ class IcebergReflectionSuite extends AnyFunSuite { assert(ex.getCause.getMessage == "boom") } + test("taskCommitFileLocations reads plain SparkWrite task commits") { + val commit = new PlainTaskCommit( + Array( + new LocationFile("s3://bucket/data/a.parquet"), + new LocationFile("s3://bucket/data/b.parquet"))) + assert( + IcebergReflection.taskCommitFileLocations(commit) == + Seq("s3://bucket/data/a.parquet", "s3://bucket/data/b.parquet")) + } + + test("taskCommitFileLocations reads new files from position-delta task commits") { + val commit = new DeltaTaskCommit( + Array(new LocationFile("s3://bucket/data/a.parquet")), + Array(new LocationFile("s3://bucket/delete/d.parquet")), + Array(new LocationFile("s3://bucket/delete/old.parquet"))) + + assert( + IcebergReflection.taskCommitFileLocations(commit) == + Seq("s3://bucket/data/a.parquet", "s3://bucket/delete/d.parquet")) + } + test("getFileFormat reads format() when declared") { val file = new FormatFile("PARQUET") assert(IcebergReflection.getFileFormat(classOf[FormatFile], file) == Some("PARQUET")) @@ -309,6 +330,19 @@ class IcebergReflectionSuite extends AnyFunSuite { assert(spec.partitionType().fields().get(0).`type`().toString == "string") } + class PlainTaskCommit(taskFiles: Array[LocationFile]) { + def files(): Array[LocationFile] = taskFiles + } + + class DeltaTaskCommit( + data: Array[LocationFile], + deletes: Array[LocationFile], + rewrittenDeletes: Array[LocationFile]) { + def dataFiles(): Array[LocationFile] = data + def deleteFiles(): Array[LocationFile] = deletes + def rewrittenDeleteFiles(): Array[LocationFile] = rewrittenDeletes + } + /** Mimics a newer Iceberg ContentFile, which exposes location(). */ class LocationFile(loc: String) { def location(): String = loc diff --git a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala index babe86b3c72..2feb4e12e2d 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterSuite.scala @@ -34,16 +34,18 @@ import org.apache.parquet.hadoop.util.HadoopInputFile import org.apache.parquet.schema.{MessageType, Type} import org.apache.spark.internal.io.FileCommitProtocol import org.apache.spark.sql.{AnalysisException, DataFrame, Row, SaveMode} -import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeScanExec, CometScanExec, CometWriteFilesExec} -import org.apache.spark.sql.execution.{FileSourceScanExec, SparkPlan} -import org.apache.spark.sql.execution.datasources.{BasicWriteTaskStats, SQLHadoopMapReduceCommitProtocol, WriteTaskStats, WriteTaskStatsTracker} +import org.apache.spark.sql.catalyst.{InternalRow, TableIdentifier} +import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeColumnarToRowExec, CometNativeScanExec, CometNativeWriteExec, CometScanExec, CometSparkToColumnarExec, CometWriteFilesExec} +import org.apache.spark.sql.execution.{ColumnarToRowTransition, FileSourceScanExec, SparkPlan, SQLExecution} +import org.apache.spark.sql.execution.command.DataWritingCommandExec +import org.apache.spark.sql.execution.datasources.{BasicWriteTaskStats, SQLHadoopMapReduceCommitProtocol, WriteFilesExec, WriteTaskStats, WriteTaskStatsTracker} import org.apache.spark.sql.functions.{array, col, map, struct, when} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, LongType, MapType, Metadata, MetadataBuilder, StringType, StructField, StructType} import org.apache.comet.CometConf import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} +import org.apache.comet.rules.RevertNativeForTransitionHeavyStages import org.apache.comet.serde.operator.NativeWriteUtils import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, SchemaGenOptions} @@ -126,6 +128,161 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } } + for (adaptive <- Seq(false, true)) { + test(s"transition-heavy fallback restores the Spark write plan with AQE=$adaptive") { + // Force a transition below the writer on both paths: Spark 3.x replaces the command, + // whereas Spark 4.x replaces only its WriteFilesExec child. + withTempPath { dir => + withTempPath { inputDir => + val inputPath = createTestData(inputDir) + val nativeOutput = new File(dir, "native-output").getAbsolutePath + val commonConf = Seq( + CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Halifax", + CometConf.COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) + + withSQLConf( + (commonConf :+ + (CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false")): _*) { + val df = spark.read.parquet(inputPath) + val nativePlan = captureWritePlan(path => df.write.parquet(path), nativeOutput) + assertHasCometNativeWriteExec(nativePlan) + verifyWrittenFile(nativeOutput) + val nativeWrite: SparkPlan = nativePlan + .collectFirst { + case write: CometNativeWriteExec => write + case write: CometWriteFilesExec => write + } + .getOrElse(fail(s"expected a native parquet writer:\n$nativePlan")) + val stageWithTransition = nativeWrite.withNewChildren(Seq( + CometSparkToColumnarExec(CometNativeColumnarToRowExec(nativeWrite.children.head)))) + assert( + stageWithTransition.collect { case _: ColumnarToRowTransition => true }.nonEmpty, + s"test requires a C2R transition:\n$stageWithTransition") + + var fallbackPlan: SparkPlan = null + withSQLConf( + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") { + val stagePlan = if (isSpark40Plus) { + val command = nativePlan + .collectFirst { case node: DataWritingCommandExec => node } + .getOrElse(fail(s"expected a Spark write command:\n$nativePlan")) + command.withNewChildren(Seq(stageWithTransition)) + } else { + stageWithTransition + } + fallbackPlan = RevertNativeForTransitionHeavyStages(spark)(stagePlan) + } + assertNoCometNativeWriteExec(fallbackPlan) + val commands = fallbackPlan.collect { case command: DataWritingCommandExec => + command + } + assert( + commands.exists(_.child.isInstanceOf[WriteFilesExec]), + s"expected DataWritingCommandExec -> WriteFilesExec after fallback:\n$fallbackPlan") + + deletePath(nativeOutput) + SQLExecution.withNewExecutionId(spark.range(0).queryExecution) { + fallbackPlan.executeCollect() + } + verifyWrittenFile(nativeOutput) + } + } + } + } + } + + for (adaptive <- Seq(false, true)) { + test( + s"transition-heavy fallback restores parquet writes through the query pipeline with AQE=$adaptive") { + withTempPath { dir => + val nativeOutput = new File(dir, "native-output").getAbsolutePath + val fallbackOutput = new File(dir, "fallback-output").getAbsolutePath + // Row source, matching the Iceberg VALUES e2e: native write engages via SparkToColumnar. + val commonConf = Seq( + CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", + CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "true", + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Halifax", + CometConf.COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) + + def rangeWrite(path: String): Unit = + spark.range(0, 1000, 1, numPartitions = 4).toDF("id").write.parquet(path) + + val sparkOutput = new File(dir, "spark-output").getAbsolutePath + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + rangeWrite(sparkOutput) + } + + withSQLConf( + (commonConf :+ + (CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false")): _*) { + val nativePlan = captureWritePlan(rangeWrite, nativeOutput) + assertHasCometNativeWriteExec(nativePlan) + assertWrittenRows(nativeOutput, sparkOutput, 1000) + } + + withSQLConf( + (commonConf ++ Seq( + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0")): _*) { + val fallbackPlan = captureWritePlan(rangeWrite, fallbackOutput) + assertTransitionHeavyParquetFallback(fallbackPlan) + assertWrittenRows(fallbackOutput, sparkOutput, 1000) + } + } + } + } + + for (adaptive <- Seq(false, true)) { + test(s"transition-heavy fallback restores a union of row sources with AQE=$adaptive") { + withTempPath { dir => + val nativeOutput = new File(dir, "native-union-output").getAbsolutePath + val fallbackOutput = new File(dir, "fallback-union-output").getAbsolutePath + val commonConf = Seq( + CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", + CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "true", + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Halifax", + CometConf.COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) + + def unionedWrite(path: String): Unit = { + val left = spark.range(0, 1000, 1, numPartitions = 2).toDF("id") + val mid = spark.range(1000, 2000, 1, numPartitions = 2).toDF("id") + val right = spark.range(2000, 3000, 1, numPartitions = 2).toDF("id") + left.union(mid).union(right).write.parquet(path) + } + + val sparkOutput = new File(dir, "spark-output").getAbsolutePath + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + unionedWrite(sparkOutput) + } + + withSQLConf( + (commonConf :+ + (CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false")): _*) { + val nativePlan = captureWritePlan(unionedWrite, nativeOutput) + assertHasCometNativeWriteExec(nativePlan) + assertWrittenRows(nativeOutput, sparkOutput, 3000) + } + + withSQLConf( + (commonConf ++ Seq( + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0")): _*) { + val fallbackPlan = captureWritePlan(unionedWrite, fallbackOutput) + assertTransitionHeavyParquetFallback(fallbackPlan) + assertWrittenRows(fallbackOutput, sparkOutput, 3000) + } + } + } + } + test("basic parquet write with repartition") { withTempPath { dir => // Create test data and write it to a temp parquet file first @@ -456,13 +613,10 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { checkAnswer(spark.read.parquet(outputPath), df.collect()) assertParquetCodec(outputPath, CompressionCodecName.GZIP) - if (isSpark40Plus) { - // Spark names the file; Comet fills it. The extension is the only externally visible - // statement of the codec, so it has to agree with the footer. (On 3.x the native writer - // invents a name with no codec suffix, so there is nothing to compare.) - listPartFileNames(outputPath).foreach { name => - assert(name.endsWith(".gz.parquet"), s"Expected a gzip file name, got '$name'") - } + // Spark names the file; Comet fills it. The extension is the only externally visible + // statement of the codec, so it has to agree with the footer. + listPartFileNames(outputPath).foreach { name => + assert(name.endsWith(".gz.parquet"), s"Expected a gzip file name, got '$name'") } } } @@ -1080,16 +1234,13 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } // --------------------------------------------------------------------------------------------- - // Spark 4.0+ only. These cover behavior that comes from leaving Spark's write framework in - // place, which is only possible where `V1WritesUtils.getWriteFilesOpt` matches the - // `WriteFilesExecBase` trait. See CometWriteFilesExec. + // Commit-protocol checks run on both writers. Tests requiring the surrounding Spark write + // command are gated to Spark 4.0+, where Comet replaces only WriteFilesExec. // --------------------------------------------------------------------------------------------- test("write creates a _SUCCESS marker") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") // https://github.com/apache/datafusion-comet/issues/2985 - the marker comes from - // HadoopMapReduceCommitProtocol.commitJob, which only runs because Comet leaves - // InsertIntoHadoopFsRelationCommand in the plan. + // HadoopMapReduceCommitProtocol.commitJob on both native writer paths. withTempPath { dir => val outputPath = new File(dir, "output.parquet").getAbsolutePath withTempPath { srcDir => @@ -1109,7 +1260,6 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } test("written file names follow Spark's naming convention") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") // The file name comes from FileCommitProtocol.newTaskTempFile and must be used verbatim: // part---c..parquet. Committers that track individual files // and tools that parse these names depend on it. @@ -1138,12 +1288,9 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } } - test("a custom output basename is honored on local storage") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") - // `escapedHdfsDestination` gates the basename on HDFS only, where the native URL parser would - // rename the file out from under the committer. Local writes hand the path to the native - // writer verbatim, so this is the control that keeps that guard from being widened: `part%foo` - // is declined on HDFS and has to keep working here. + test("local output basenames match Spark's version-specific commit protocol") { + // Spark 4.0+ honors the basename, including literal '%' on local storage where the path + // reaches the native writer verbatim. Spark 3.x ignores the option and always uses "part". Seq("out", "part%foo").foreach { basename => withTempPath { dir => val outputPath = new File(dir, "output.parquet").getAbsolutePath @@ -1159,9 +1306,10 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } val written = new File(outputPath).listFiles().map(_.getName).filter(_.endsWith(".parquet")) + val expectedBasename = if (isSpark40Plus) basename else "part" assert( - written.nonEmpty && written.forall(_.startsWith(s"$basename-")), - s"expected every data file to be named '$basename-...', found: " + + written.nonEmpty && written.forall(_.startsWith(s"$expectedBasename-")), + s"expected every data file to be named '$expectedBasename-...', found: " + written.mkString(", ")) checkAnswer(spark.read.parquet(outputPath), df) } @@ -1170,9 +1318,9 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } test("INSERT INTO ... SELECT is visible to subsequent reads") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") // https://github.com/apache/datafusion-comet/issues/3521 - reads returned no rows because the - // bespoke write path never refreshed the catalog cache. Spark's command does that itself. + // bespoke write path never refreshed the catalog cache. On Spark 4.0+ Spark's command does + // that itself; on 3.x CometNativeWriteExec repeats it after the job commits. withTable("comet_write_target", "comet_write_source") { withNativeWriter { sql("CREATE TABLE comet_write_source(id bigint, name string) USING parquet") @@ -1192,9 +1340,127 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } } + test("a native overwrite of a cached path recaches the data") { + withTempPath { dir => + val target = new File(dir, "target").getAbsolutePath + val source = new File(dir, "source").getAbsolutePath + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark.range(0, 5).write.parquet(target) + spark.range(100, 103).write.parquet(source) + } + val cached = spark.read.parquet(target).cache() + try { + assert(cached.count() == 5) + withNativeWriter { + assertHasCometNativeWriteExec( + captureWritePlan( + spark.read.parquet(source).write.mode(SaveMode.Overwrite).parquet(target))) + } + // Spark recaches every cached plan that reads the output path once the write commits. + checkAnswer(cached, (100L until 103L).map(Row(_))) + } finally { + cached.unpersist() + } + } + } + + test("a native write rejects case-insensitively duplicate column names") { + withTempPath { dir => + val source = new File(dir, "source").getAbsolutePath + val target = new File(dir, "target") + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark.range(3).selectExpr("id AS a", "id AS b").write.parquet(source) + } + withNativeWriter { + val (plan, e) = captureFailedWritePlan( + spark.read + .parquet(source) + .selectExpr("a", "b AS A") + .write + .mode(SaveMode.Overwrite) + .parquet(target.getAbsolutePath)) + // Spark's own writer rejects these columns too, so a fallback would also pass the checks + // below without exercising the native writer's check. + assertHasCometNativeWriteExec(plan) + assert(e.isInstanceOf[AnalysisException], e) + assert(e.getMessage.contains("COLUMN_ALREADY_EXISTS"), e.getMessage) + } + assert(!target.exists(), "nothing may be written when the columns are rejected") + } + } + + test("a native INSERT into a cached table recaches the table") { + withTable("comet_write_target", "comet_write_source") { + sql("CREATE TABLE comet_write_source(id bigint) USING parquet") + sql("CREATE TABLE comet_write_target(id bigint) USING parquet") + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sql("INSERT INTO comet_write_source VALUES (1), (2)") + sql("INSERT INTO comet_write_target VALUES (0)") + } + sql("CACHE TABLE comet_write_target") + try { + checkAnswer(spark.table("comet_write_target"), Row(0L)) + withNativeWriter { + assertHasCometNativeWriteExec( + captureWritePlan( + sql("INSERT INTO comet_write_target SELECT id FROM comet_write_source"))) + } + checkAnswer(spark.table("comet_write_target"), Seq(Row(0L), Row(1L), Row(2L))) + } finally { + sql("UNCACHE TABLE IF EXISTS comet_write_target") + } + } + } + + for { + autoUpdate <- Seq(true, false) + overwrite <- Seq(false, true) + } { + val insert = if (overwrite) "INSERT OVERWRITE" else "INSERT INTO" + test(s"native $insert maintains table stats with size autoUpdate=$autoUpdate") { + withTable("comet_write_target", "comet_write_source") { + withSQLConf(SQLConf.AUTO_SIZE_UPDATE_ENABLED.key -> autoUpdate.toString) { + sql("CREATE TABLE comet_write_source(id bigint) USING parquet") + sql("CREATE TABLE comet_write_target(id bigint) USING parquet") + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + sql("INSERT INTO comet_write_source SELECT id FROM range(100)") + sql("INSERT INTO comet_write_target VALUES (0)") + } + sql("ANALYZE TABLE comet_write_target COMPUTE STATISTICS NOSCAN") + val target = TableIdentifier("comet_write_target") + val catalog = spark.sessionState.catalog + val statsBefore = catalog.getTableMetadata(target).stats + assert(statsBefore.isDefined) + + withNativeWriter { + assertHasCometNativeWriteExec( + captureWritePlan( + sql(s"$insert comet_write_target SELECT id FROM comet_write_source"))) + } + + // Mirrors CommandUtils.updateTableStats, which Spark runs after the write commits. + val statsAfter = catalog.getTableMetadata(target).stats + if (autoUpdate) { + val location = new Path(catalog.getTableMetadata(target).location) + val fs = location.getFileSystem(spark.sessionState.newHadoopConf()) + val dataSize = fs + .listStatus(location) + .filterNot(f => + f.getPath.getName.startsWith("_") || f.getPath.getName.startsWith(".")) + .map(_.getLen) + .sum + assert(statsAfter.map(_.sizeInBytes) == Some(BigInt(dataSize))) + assert(statsAfter != statsBefore) + } else { + assert(statsAfter.isEmpty, s"stale stats were left in place: $statsAfter") + } + } + } + } + } + test("dynamic partition overwrite falls back to Spark") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") - // A dynamic overwrite is a partitioned write, which CometWriteFiles declines - but the + // A dynamic overwrite is a partitioned write, which both native writers decline - but the // consequence of getting it wrong is silent data loss across untouched partitions, so assert // the fallback and the semantics explicitly rather than relying on the partitioning check. withTempPath { dir => @@ -1254,14 +1520,14 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } test("empty input still writes a schema-only file (SPARK-23271)") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") // An empty input must still leave a schema behind for downstream readers: `spark.read.parquet` // of the output must see the write's schema, not fail. Comet reaches this in two ways - if the - // native child has one partition producing no batches, the partition-0 branch of executeTask - // writes a metadata-only file; if it produces zero partitions, doExecuteWrite swaps in a dummy - // single-partition RDD to get to the same branch. This test exercises the first; the - // zero-partition swap is reached by an AQE-collapsed empty relation and is covered by - // CometEmptyRelationParquetWriterSuite. + // native child has one partition producing no batches, the partition-0 branch of the write + // task writes a metadata-only file; if it produces zero partitions, the writer swaps in a + // dummy single-partition RDD to get to the same branch. Both CometWriteFilesExec (Spark 4.0+) + // and CometNativeWriteExec (Spark 3.x) do this. This test exercises the first; the + // zero-partition swap is covered by the empty-directory test below and, on Spark 4.0+, + // by CometEmptyRelationParquetWriterSuite. withTempPath { dir => val outputPath = new File(dir, "output.parquet").getAbsolutePath val sourcePath = new File(dir, "source.parquet").getAbsolutePath @@ -1282,6 +1548,31 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } } + test("a zero-partition native scan writes a schema-only file (#5303)") { + withTempPath { dir => + val source = new File(dir, "empty-source") + assert(source.mkdirs()) + val output = new File(dir, "output") + withNativeWriter { + val empty = spark.read.schema("id INT, name STRING").parquet(source.getAbsolutePath) + val scans = collect(empty.queryExecution.executedPlan) { case scan: CometNativeScanExec => + scan + } + assert(scans.size == 1, empty.queryExecution.executedPlan.toString) + assert(scans.head.executeColumnar().getNumPartitions == 0) + + val plan = captureWritePlan(p => empty.write.parquet(p), output.getAbsolutePath) + assertHasCometNativeWriteExec(plan) + assert(new File(output, "_SUCCESS").isFile) + assert(listPartFileNames(output.getAbsolutePath).size == 1) + assert(!new File(output, "_temporary").exists()) + val readBack = spark.read.parquet(output.getAbsolutePath) + assert(readBack.schema == empty.schema) + assert(readBack.collect().isEmpty) + } + } + } + test("INSERT INTO ... SELECT writes the target table's column names") { assume(isSpark40Plus, "Requires the WriteFilesExec seam") // https://github.com/apache/datafusion-comet/issues/3426, which is Spark's own @@ -1330,8 +1621,7 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } test("an empty partition writes no file and still commits") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") - // executeTask's `sparkPartitionId != 0 && !batches.hasNext` branch must skip newTaskTempFile + // The write task's `partitionId != 0 && !batches.hasNext` branch must skip newTaskTempFile // altogether and still commit the task, matching FileFormatWriter's EmptyDirectoryDataWriter. // Hash-partitioning into eight and keeping a single id leaves at most one partition with // rows, so at most two files can appear: that one, plus partition 0's schema-only file when @@ -1390,8 +1680,7 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } test("a failing task aborts and cleans up its staging file") { - assume(isSpark40Plus, "Requires the WriteFilesExec seam") - // CometWriteFilesExec.executeTask must call committer.abortTask and rethrow. Injecting the + // Both native writers must call committer.abortTask and rethrow. Injecting the // failure through the commit protocol rather than the data lets the write get as far as // creating a staging file, so the cleanup is actually observable. // @@ -1442,6 +1731,41 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } } + test("a row group the pool cannot hold fails its task with a native out-of-memory error") { + // The native writer reserves its in-progress row group from the task's memory pool, so a row + // group that outgrows the pool fails the task, which Spark can retry, instead of growing in + // memory that no budget accounts for until the executor is killed. + withTempPath { dir => + val outputPath = new File(dir, "output.parquet").getAbsolutePath + val sourcePath = new File(dir, "source.parquet").getAbsolutePath + withNativeWriter { + // One task writes about 10 MB of random strings, which snappy cannot shrink, into a single + // row group, against a pool of about 4 MiB: 0.002 of the suite's 2 GiB off-heap size. + val rows = 20000 + val df = materializeAsCometSource( + (0 until rows) + .map(i => (i, new Random(i).alphanumeric.take(500).mkString)) + .toDF("id", "payload") + .repartition(1), + sourcePath) + withSQLConf(CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> "0.002") { + val e = intercept[Exception] { + df.write.parquet(outputPath) + } + val outOfMemory = causeChain(e) + .flatMap(t => Option(t.getMessage)) + .exists(_.contains("Additional allocation failed for ParquetWriterExec")) + assert(outOfMemory, s"Expected the native writer's out-of-memory error, got: $e") + } + + // The executor outlived the failure, and with the whole pool the same write succeeds. + val plan = captureWritePlan(p => df.write.mode(SaveMode.Overwrite).parquet(p), outputPath) + assertHasCometNativeWriteExec(plan) + assert(spark.read.parquet(outputPath).count() == rows) + } + } + } + test("the Spark 3.x opt-in key still enables native writes on Spark 4.0+") { assume(isSpark40Plus, "Requires the WriteFilesExec seam") // The opt-in moved from DataWritingCommandExec to WriteFilesExec with the operator. Jobs that @@ -1514,6 +1838,58 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { } } + private def deletePath(path: String): Unit = { + def delete(file: File): Unit = { + if (file.isDirectory) { + Option(file.listFiles()).foreach(_.foreach(delete)) + } + file.delete() + } + delete(new File(path)) + } + + /** + * Spark 3.x replaces the whole command with `CometNativeWriteExec`, and Spark then inserts a + * C2R above that columnar node, so `maxTransitions = 0` reverts the stage back to + * `DataWritingCommandExec` -> `WriteFilesExec`. Spark 4.0+ leaves the command in place and does + * not insert a C2R above `CometWriteFilesExec`. These row sources only add `SparkToColumnar`, + * which is an R2C, so the stage is not transition-heavy and the native write stays. + */ + private def assertTransitionHeavyParquetFallback(plan: SparkPlan): Unit = { + if (isSpark40Plus) { + assertHasCometNativeWriteExec(plan) + } else { + assertNoCometNativeWriteExec(plan) + assertRestoredParquetWriteCommand(plan) + } + } + + private def assertRestoredParquetWriteCommand(plan: SparkPlan): Unit = { + assert( + plan + .collect { case command: DataWritingCommandExec => command } + .exists(_.child.isInstanceOf[WriteFilesExec]), + s"expected DataWritingCommandExec -> WriteFilesExec after fallback:\n$plan") + } + + // Compare every row with Spark-written data using the row-based parquet reader, so this + // check is independent of Comet scans and detects duplicates or missing union inputs. + private def assertWrittenRows( + outputPath: String, + sparkOutputPath: String, + expectedRows: Int): Unit = { + val outputDir = new File(outputPath) + val partFiles = outputDir.listFiles().filter(_.getName.startsWith("part-")) + assert(partFiles.length > 1, s"Expected multiple part files under $outputPath") + withSQLConf( + CometConf.COMET_ENABLED.key -> "false", + SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "false") { + val expected = spark.read.parquet(sparkOutputPath).collect().toSeq + assert(expected.size == expectedRows) + checkAnswer(spark.read.parquet(outputPath), expected) + } + } + private def writeWithCometNativeWriteExec( inputPath: String, outputPath: String, @@ -1529,10 +1905,10 @@ class CometParquetWriterSuite extends CometParquetWriterTestBase { Some(plan) } - private def verifyWrittenFile(outputPath: String): Unit = { + private def verifyWrittenFile(outputPath: String, expectedRows: Int = 1000): Unit = { // Verify the data was written correctly val resultDf = spark.read.parquet(outputPath) - assert(resultDf.count() == 1000, "Expected 1000 rows to be written") + assert(resultDf.count() == expectedRows, s"Expected $expectedRows rows to be written") // Verify multiple part files were created val outputDir = new File(outputPath) diff --git a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala index 20caba36b1f..605dcbbb3ee 100644 --- a/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala +++ b/spark/src/test/scala/org/apache/comet/parquet/CometParquetWriterTestBase.scala @@ -21,6 +21,8 @@ package org.apache.comet.parquet import java.util.concurrent.atomic.AtomicReference +import scala.util.control.NonFatal + import org.apache.spark.CometListenerBusUtils import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.comet.{CometNativeWriteExec, CometWriteFilesExec} @@ -64,19 +66,37 @@ abstract class CometParquetWriterTestBase extends CometTestBase { /** As above, for a write that names its own target (an `INSERT INTO`, for example). */ protected def captureWritePlan(writeOp: => Unit): SparkPlan = { + val (plan, failure) = captureWriteOutcome(writeOp) + failure.foreach(throw _) + plan.getOrElse(fail("Listener was not called - no execution plan captured")) + } + + /** + * As above, for a write that is expected to fail. Returns the plan of the failed write, so a + * test can check that the failure came from the native writer rather than from Spark's. + */ + protected def captureFailedWritePlan(writeOp: => Unit): (SparkPlan, Throwable) = { + val (plan, failure) = captureWriteOutcome(writeOp) + val error = failure.getOrElse(fail("Expected the write to fail, but it succeeded")) + (plan.getOrElse(fail("Listener was not called - no execution plan captured")), error) + } + + private def captureWriteOutcome(writeOp: => Unit): (Option[SparkPlan], Option[Throwable]) = { val capturedPlan = new AtomicReference[QueryExecution]() + def isWrite(funcName: String): Boolean = funcName == "save" || funcName.contains("command") val listener = new org.apache.spark.sql.util.QueryExecutionListener { override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = { - if (funcName == "save" || funcName.contains("command")) { + if (isWrite(funcName)) { capturedPlan.set(qe) } } - override def onFailure( - funcName: String, - qe: QueryExecution, - exception: Exception): Unit = {} + override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = { + if (isWrite(funcName)) { + capturedPlan.set(qe) + } + } } // Listener events are delivered asynchronously, so drain the bus before registering: an @@ -85,12 +105,15 @@ abstract class CometParquetWriterTestBase extends CometTestBase { spark.listenerManager.register(listener) try { - writeOp + val failure = + try { + writeOp + None + } catch { + case NonFatal(e) => Some(e) + } CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) - - val plan = capturedPlan.get() - assert(plan != null, "Listener was not called - no execution plan captured") - stripAQEPlan(plan.executedPlan) + (Option(capturedPlan.get()).map(qe => stripAQEPlan(qe.executedPlan)), failure) } finally { spark.listenerManager.unregister(listener) } diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala index 006cc1ae126..d8f40f2b635 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -2356,6 +2356,20 @@ class CometExecRuleSuite extends CometTestBase { } } + test("only the OneRowRelation conversion is on by default") { + withConversionQueries { queries => + // CometTestBase turns every conversion on, so set each back to its default. + val defaults = queries.map { case (entry, _) => entry.key -> entry.defaultValueString } + for ((entry, query) <- queries) { + withClue(s"$query: ") { + assert( + convertedLeaves(query, defaults: _*).nonEmpty == + (entry eq CometConf.COMET_CONVERT_FROM_ONE_ROW_RELATION_ENABLED)) + } + } + } + } + test("the deprecated sparkToColumnar settings still convert what they converted before") { val switch = CometConf.COMET_SPARK_TO_ARROW_ENABLED.key -> "true" val list = CometConf.COMET_SPARK_TO_ARROW_SUPPORTED_OPERATOR_LIST.key diff --git a/spark/src/test/scala/org/apache/comet/rules/CometScanSchemeFallbackSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometScanSchemeFallbackSuite.scala index 5a14f9f592e..21011fae65d 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometScanSchemeFallbackSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometScanSchemeFallbackSuite.scala @@ -27,11 +27,12 @@ import java.util.UUID import org.apache.commons.io.FileUtils import org.apache.spark.SparkConf import org.apache.spark.sql.{CometTestBase, SaveMode} -import org.apache.spark.sql.comet.CometScanExec +import org.apache.spark.sql.comet.{CometIcebergNativeScanExec, CometIcebergWriteExec, CometScanExec} import org.apache.spark.sql.execution.{FileSourceScanExec, SparkPlan} -import org.apache.comet.CometConf -import org.apache.comet.hadoop.fs.{FakeHDFSFileSystem, FakeHdfsSchemeFileSystem} +import org.apache.comet.{CometConf, CometIcebergTestBase, ExtendedExplainInfo, NativeBase} +import org.apache.comet.hadoop.fs.{FakeHDFSFileSystem, FakeHdfsSchemeFileSystem, FakeWasbSchemeFileSystem} +import org.apache.comet.iceberg.IcebergStorageSchemes /** * Comet's native readers go through object_store, which understands a fixed set of URL schemes. A @@ -43,7 +44,7 @@ import org.apache.comet.hadoop.fs.{FakeHDFSFileSystem, FakeHdfsSchemeFileSystem} * native probe is cached per probe URL so an authorityless URL can't poison the authority-bearing * form of the same scheme. */ -class CometScanSchemeFallbackSuite extends CometTestBase { +class CometScanSchemeFallbackSuite extends CometTestBase with CometIcebergTestBase { private var fakeRootDir: File = _ @@ -54,6 +55,9 @@ class CometScanSchemeFallbackSuite extends CometTestBase { // cluster. `hdfs` is natively readable by default, so this scan must be CLAIMED, not declined. conf.set("spark.hadoop.fs.hdfs.impl", "org.apache.comet.hadoop.fs.FakeHdfsSchemeFileSystem") conf.set("spark.hadoop.fs.defaultFS", FakeHDFSFileSystem.PREFIX) + // Back `wasb` with a local FS so an Iceberg table can live under a `wasb://` warehouse. The + // native Iceberg storage factory has no arm for it, so that scan must be DECLINED, not claimed. + conf.set("spark.hadoop.fs.wasb.impl", classOf[FakeWasbSchemeFileSystem].getName) // Intentionally NOT setting CometConf.COMET_LIBHDFS_SCHEMES -- `fake` is not natively readable, // and `hdfs` must still be claimed by default (mirrors the native `is_hdfs_scheme` default). conf @@ -120,10 +124,65 @@ class CometScanSchemeFallbackSuite extends CometTestBase { "gs://bucket/... must be admitted even after an authorityless gs:// URI was probed first") } + test("iceberg gate: the JNI scheme list parser trims, lowercases and rejects a blank list") { + assert( + IcebergStorageSchemes.parse("file, S3,,s3a,", forWrite = false) == Set("file", "s3", "s3a")) + // A loaded library publishes a fixed non-empty list, so a blank answer is a build bug and + // must fail loudly rather than quietly disable every native Iceberg scan and write. + Seq(("", false), (" , ", true), (null, false)).foreach { case (joined, forWrite) => + val path = if (forWrite) "write" else "read" + val e = intercept[IllegalStateException](IcebergStorageSchemes.parse(joined, forWrite)) + assert( + e.getMessage == s"Comet native library published no Iceberg $path scheme list", + s"unexpected message for ${Option(joined)}: ${e.getMessage}") + } + } + + test("iceberg gate: an unloaded native library yields no schemes") { + // Every caller sits behind `isCometLoaded`, so this answer never reaches a plan; it only + // guarantees that nothing hand-maintained stands in for the native list. + assert(IcebergStorageSchemes.load(forWrite = false, isLoaded = false) == Set.empty) + assert(IcebergStorageSchemes.load(forWrite = true, isLoaded = false) == Set.empty) + } + + test("iceberg gate: the scheme sets come from native and carry its mode-specific entries") { + // The lazy sets must be what the JNI probe answers, not an empty set from a load that ran + // before the library was ready. `oss` is read-only (no `oss.*` property forwarding for + // writes) and `memory` is write-only (a fresh in-process store per FileIO, so a read finds + // nothing); the Azure schemes and `gcs` have no storage factory arm at all. + assume(NativeBase.isLoaded, "Comet native library not loaded") + val read = IcebergStorageSchemes.read + val write = IcebergStorageSchemes.write + assert(read == IcebergStorageSchemes.parse(NativeBase.icebergStorageSchemes(false), false)) + assert(write == IcebergStorageSchemes.parse(NativeBase.icebergStorageSchemes(true), true)) + assert(read.nonEmpty, "native published no read schemes") + assert(write.nonEmpty, "native published no write schemes") + assert(read.contains("oss"), s"oss must be in the native read schemes ${read.toSeq.sorted}") + assert( + !write.contains("oss"), + s"oss must not be in the native write schemes ${write.toSeq.sorted}") + assert( + write.contains("memory"), + s"memory must be in the native write schemes ${write.toSeq.sorted}") + assert( + !read.contains("memory"), + s"memory must not be in the native read schemes ${read.toSeq.sorted}") + Seq("abfs", "abfss", "wasb", "wasbs", "gcs").foreach { scheme => + assert( + !read.contains(scheme), + s"$scheme must not be in the native read schemes ${read.toSeq.sorted}") + assert( + !write.contains(scheme), + s"$scheme must not be in the native write schemes ${write.toSeq.sorted}") + } + } + test("iceberg gate: builtin allowlist admitted, unbuildable schemes rejected") { - // The Iceberg gate is an explicit allowlist mirroring `storage_factory_for`'s arms (keep in - // lockstep). Narrower than the Parquet gate: object_store recognizes http/abfs/wasb but - // iceberg-rust can't build them, so reject up-front rather than fail during native setup. + // The Iceberg gate is the allowlist the native storage factory publishes over JNI. Narrower + // than the Parquet gate: object_store recognizes http and the Azure schemes, but the native + // Iceberg storage factory has no arm for them, so reject up-front rather than fail at + // execution. `memory` is write-only (a fresh in-process store per FileIO holds no table), and + // the built-in set matches verbatim because native opens a location by its raw scheme. Seq( "file:///tmp/key.parquet", "s3://bucket/key.parquet", @@ -135,6 +194,9 @@ class CometScanSchemeFallbackSuite extends CometTestBase { s"$u must be iceberg-readable; icebergReadableSchemes has regressed") } Seq( + "memory:///key.parquet", + "S3://bucket/key.parquet", + "File:///tmp/key.parquet", "http://bucket.example.com/key.parquet", "https://bucket.example.com/key.parquet", "abfs://container@acct/key.parquet", @@ -143,10 +205,31 @@ class CometScanSchemeFallbackSuite extends CometTestBase { "wasbs://container@acct/key.parquet").foreach { u => assert( !CometScanRule.isIcebergReadableScheme(new URI(u), Set.empty), - s"$u must not be iceberg-readable; storage_factory_for has no matching arm") + s"$u must not be iceberg-readable; storage_factory_for rejects it for reads") } } + test("iceberg gate: an opt-in alias is matched verbatim, like the built-in schemes") { + // Native opens an alias location as written and OpenDAL's S3 backend checks it against a + // lowercase `scheme://bucket/` prefix, so a `BLOB://` location admitted here would only fail + // at execution. The alias set is lowercase, so a scheme written any other way is declined. + val schemes = Set("blob") + assert( + CometScanRule.isIcebergReadableScheme(new URI("blob://bucket/key.parquet"), schemes), + "blob://bucket/... must be admitted once blob is opted in") + Seq("BLOB://bucket/key.parquet", "Blob://bucket/key.parquet", "BLOB:///bucket/key.parquet") + .foreach { u => + assert( + !CometScanRule.isIcebergReadableScheme(new URI(u), schemes), + s"$u must be declined: native opens the location as written and the S3 backend " + + "rejects a scheme prefix that is not lowercase") + } + // The fallback reason names the lowercase alias, which is the form the scan would admit. + assert( + CometScanRule.icebergSupportedSchemesMessage(schemes).contains("blob"), + "the supported-schemes message must list the opted-in alias in its admitted form") + } + test("parquet gate: mixed-bucket alias scan is declined (single object store per partition)") { // Native planning registers one object store per FilePartition and strips the authority from // every file's object key, so files in a second bucket would be read from the first. A scan @@ -218,6 +301,67 @@ class CometScanSchemeFallbackSuite extends CometTestBase { "without opt-in, blob gets no bucket promotion, so a hostless blob location is unopenable") } + test("iceberg scan and write decline a wasb table instead of failing at execution") { + // Regression guard for both scheme gates: the native Iceberg storage factory has no arm for + // `wasb`, so the scan and the write must each be declined with a reason naming the scheme + // and run in Spark, rather than be claimed and fail at execution with "Unsupported storage + // scheme: wasb". + assume(icebergAvailable, "Iceberg not available in classpath") + withTempIcebergDir { dir => + val warehouse = s"${FakeWasbSchemeFileSystem.PREFIX}${dir.getAbsolutePath}/warehouse" + withSQLConf( + "spark.sql.catalog.wasb_catalog" -> "org.apache.iceberg.spark.SparkCatalog", + "spark.sql.catalog.wasb_catalog.type" -> "hadoop", + "spark.sql.catalog.wasb_catalog.warehouse" -> warehouse, + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true", + CometConf.COMET_ICEBERG_NATIVE_WRITE_ENABLED.key -> "true", + CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "true", + // The VALUES source must be native, or the write is declined for its input before the + // write serde (and its scheme gate) is ever consulted. + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + spark.sql("CREATE TABLE wasb_catalog.db.wasb_table (id INT, name STRING) USING iceberg") + try { + val insertPlans = capturePlans(spark) { + spark.sql( + "INSERT INTO wasb_catalog.db.wasb_table VALUES (1, 'Alice'), (2, 'Bob'), (3, 'Carol')") + } + assert(insertPlans.nonEmpty, "no INSERT plan was captured") + val nativeWrites = insertPlans.flatMap(_.collectWithSubqueries { + case w: CometIcebergWriteExec => w + }) + assert( + nativeWrites.isEmpty, + "`wasb://` has no native Iceberg storage factory arm; the write must fall back to " + + s"Spark, but Comet claimed it:\n${insertPlans.mkString("\n--\n")}") + val writeReasons = insertPlans.flatMap(new ExtendedExplainInfo().getFallbackReasons) + assert( + writeReasons.exists(_.contains("unsupported storage scheme: wasb")), + s"the write must be declined for its scheme, got reasons: $writeReasons") + + val (_, cometPlan) = checkSparkAnswerAndFallbackReason( + "SELECT * FROM wasb_catalog.db.wasb_table ORDER BY id", + "wasb") + val nativeScans = cometPlan.collect { case s: CometIcebergNativeScanExec => s } + assert( + nativeScans.isEmpty, + "`wasb://` has no native Iceberg storage factory arm; the scan must fall back to " + + s"Spark, but Comet claimed it:\n$cometPlan") + // Nothing but the scheme may have caused the fallback: every reason names `wasb`, + // except the one the operator above records about its now-Spark BatchScanExec child. + val scanReasons = new ExtendedExplainInfo().getFallbackReasons(cometPlan) + val unrelated = scanReasons.filterNot(_.contains("wasb")) + assert( + unrelated.forall(_.contains("BatchScanExec")), + s"fallback reasons other than the scheme were recorded: $unrelated") + } finally { + spark.sql("DROP TABLE wasb_catalog.db.wasb_table") + } + } + } + } + test("native scan claims hdfs:// when libhdfs.schemes is unset (native-default lockstep)") { // Native `is_hdfs_scheme` treats `hdfs` as readable when `fs.comet.libhdfs.schemes` is unset, // so the JVM gate must agree and CLAIM `hdfs://`. Guards the `case None => Set("hdfs")` default diff --git a/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala b/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala index e350badf8da..40281ce4ea1 100644 --- a/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStagesSuite.scala @@ -19,19 +19,125 @@ package org.apache.comet.rules -import org.apache.spark.sql.{CometTestBase, Row} +import java.util.concurrent.{CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.AtomicReference + +import org.apache.spark.sql.{CometTestBase, Row, SaveMode} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Literal} import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial} import org.apache.spark.sql.comet._ import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec import org.apache.spark.sql.execution._ -import org.apache.spark.sql.execution.adaptive.QueryStageExec +import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, QueryStageExec, ShuffleQueryStageExec} +import org.apache.spark.sql.execution.command.DataWritingCommandExec +import org.apache.spark.sql.execution.datasources.WriteFilesExec import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.BinaryType +import org.apache.spark.sql.util.QueryExecutionListener import org.apache.comet.CometConf import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus +import org.apache.comet.serde.OperatorOuterClass.Operator + +private case class AliasingFallbackCometExec( + override val originalPlan: SparkPlan, + child: SparkPlan) + extends CometExec + with UnaryExecNode { + override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = + copy(child = newChild) +} class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase { + private def cometIcebergWrite(child: SparkPlan): CometIcebergWriteExec = { + val output = Seq( + AttributeReference(IcebergWriteExec.CommitMessageColumn, BinaryType, nullable = false)()) + val originalPlan = IcebergWriteExec(null, output, child) + CometIcebergWriteExec( + Operator.newBuilder().build(), + originalPlan, + child, + output, + batchWrite = null, + table = null, + partitionSpecId = 0) + } + + private def cometFilter(child: SparkPlan): CometFilterExec = { + val condition = Literal.TrueLiteral + val sparkFilter = FilterExec(condition, child) + CometFilterExec( + Operator.newBuilder().build(), + sparkFilter, + sparkFilter.output, + condition, + child, + SerializedPlan(None)) + } + + private def captureDataWritingCommand(path: String): DataWritingCommandExec = { + val captured = new AtomicReference[SparkPlan]() + val callbackCompleted = new CountDownLatch(1) + val listener = new QueryExecutionListener { + override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = { + if (funcName == "save" || funcName.contains("command")) { + captured.set(qe.executedPlan) + callbackCompleted.countDown() + } + } + override def onFailure( + funcName: String, + qe: QueryExecution, + exception: Exception): Unit = {} + } + spark.listenerManager.register(listener) + try { + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark.range(1).toDF("id").write.mode("overwrite").parquet(path) + } + assert( + callbackCompleted.await(10, TimeUnit.SECONDS), + "timed out waiting to capture the parquet write plan") + } finally { + spark.listenerManager.unregister(listener) + } + val plan = stripAQEPlan( + Option(captured.get()).getOrElse(fail("expected a captured parquet write plan"))) + plan + .collectFirst { case command: DataWritingCommandExec => command } + .getOrElse(fail(s"expected DataWritingCommandExec:\n$plan")) + } + + private def cometNativeWrite( + child: SparkPlan, + command: DataWritingCommandExec): CometNativeWriteExec = { + CometNativeWriteExec( + Operator.newBuilder().build(), + command, + child, + outputPath = "/tmp/unused-native-write", + mode = SaveMode.Overwrite, + committer = null, + serializableHadoopConf = null, + outputWriterFactory = null) + } + + private def assertRestoredParquetWrite(reverted: SparkPlan): WriteFilesExec = { + val command = reverted match { + case node: DataWritingCommandExec => node + case other => fail(s"expected DataWritingCommandExec, got:\n$other") + } + val writeFiles = command.child match { + case node: WriteFilesExec => node + case other => fail(s"expected WriteFilesExec under DataWritingCommandExec, got:\n$other") + } + assert( + command.collect { case _: CometNativeWriteExec => true }.isEmpty, + s"native parquet write should be restored, not erased:\n$reverted") + writeFiles + } + private def createSparkPlan(sql: String): SparkPlan = { var plan: SparkPlan = null withSQLConf(CometConf.COMET_ENABLED.key -> "false") { @@ -59,6 +165,12 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase { plan.collect { case _: ColumnarToRowTransition => true }.size } + private def unwrapCodegen(plan: SparkPlan): SparkPlan = plan match { + case wholeStage: WholeStageCodegenExec => unwrapCodegen(wholeStage.child) + case inputAdapter: InputAdapter => unwrapCodegen(inputAdapter.child) + case other => other + } + private def collectCometAggregates(plan: SparkPlan): Seq[CometHashAggregateExec] = { val current = plan match { case aggregate: CometHashAggregateExec => Seq(aggregate) @@ -142,6 +254,207 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase { } } + test("revertToSpark preserves an Iceberg write with a leaf child") { + val sparkPlan = createSparkPlan("SELECT id FROM VALUES (1) AS t(id)") + val leaf = sparkPlan.collectFirst { case node: LeafExecNode => node }.getOrElse { + fail(s"expected a leaf node in test plan:\n$sparkPlan") + } + val write = cometIcebergWrite(leaf) + + val reverted = RevertNativeForTransitionHeavyStages(spark).revertToSpark(write) + val icebergWrite = reverted match { + case node: IcebergWriteExec => node + case other => fail(s"expected IcebergWriteExec, got:\n$other") + } + + assert(icebergWrite.child eq leaf) + assert(icebergWrite.output.map(_.name) == Seq(IcebergWriteExec.CommitMessageColumn)) + assert(icebergWrite.output.map(_.dataType) == Seq(BinaryType)) + } + + test("revertToSpark preserves an Iceberg write over SparkToColumnar of a row source") { + val sparkPlan = createSparkPlan("SELECT id FROM VALUES (1) AS t(id)") + val leaf = sparkPlan.collectFirst { case node: LeafExecNode => node }.getOrElse { + fail(s"expected a leaf node in test plan:\n$sparkPlan") + } + val write = cometIcebergWrite(CometSparkToColumnarExec(leaf)) + + val reverted = RevertNativeForTransitionHeavyStages(spark).revertToSpark(write) + val icebergWrite = reverted match { + case node: IcebergWriteExec => node + case other => fail(s"expected IcebergWriteExec, got:\n$other") + } + + assert( + icebergWrite.child eq leaf, + s"SparkToColumnar should unwrap to the row source:\n$reverted") + assert(icebergWrite.output.map(_.name) == Seq(IcebergWriteExec.CommitMessageColumn)) + assert(icebergWrite.output.map(_.dataType) == Seq(BinaryType)) + assert( + reverted.collect { case _: CometSparkToColumnarExec => true }.isEmpty, + s"SparkToColumnar should be fully unwrapped:\n$reverted") + } + + test("revertToSpark preserves an Iceberg write without duplicating its unary child") { + val sparkPlan = createSparkPlan("SELECT id FROM VALUES (1), (2) AS t(id)") + val leaf = sparkPlan.collectFirst { case node: LeafExecNode => node }.getOrElse { + fail(s"expected a leaf node in test plan:\n$sparkPlan") + } + val write = cometIcebergWrite(cometFilter(leaf)) + + val reverted = RevertNativeForTransitionHeavyStages(spark).revertToSpark(write) + + assert(reverted.isInstanceOf[IcebergWriteExec], s"expected IcebergWriteExec:\n$reverted") + assert( + reverted.collect { case _: FilterExec => true }.size == 1, + s"expected exactly one Spark FilterExec:\n$reverted") + assert(countCometExecs(reverted) == 0, s"expected no Comet operators:\n$reverted") + } + + test("revertToSpark unwraps stacked SparkToColumnar(C2R) under a native write") { + val sparkPlan = createSparkPlan("SELECT id FROM VALUES (1) AS t(id)") + val leaf = sparkPlan.collectFirst { case node: LeafExecNode => node }.getOrElse { + fail(s"expected a leaf node in test plan:\n$sparkPlan") + } + val stacked = CometSparkToColumnarExec(CometNativeColumnarToRowExec(cometFilter(leaf))) + val write = cometIcebergWrite(stacked) + + val reverted = RevertNativeForTransitionHeavyStages(spark).revertToSpark(write) + + assert(reverted.isInstanceOf[IcebergWriteExec], s"expected IcebergWriteExec:\n$reverted") + assert( + reverted.collect { case _: FilterExec => true }.size == 1, + s"expected exactly one Spark FilterExec:\n$reverted") + assert( + reverted.collect { case _: CometNativeColumnarToRowExec | _: CometSparkToColumnarExec => + true + }.isEmpty, + s"stacked transitions should be fully unwrapped:\n$reverted") + assert(countCometExecs(reverted) == 0, s"expected no Comet operators:\n$reverted") + } + + test("revertToSpark preserves a native parquet write with a leaf child") { + val sparkPlan = createSparkPlan("SELECT id FROM VALUES (1) AS t(id)") + val leaf = sparkPlan.collectFirst { case node: LeafExecNode => node }.getOrElse { + fail(s"expected a leaf node in test plan:\n$sparkPlan") + } + withTempPath { dir => + val write = cometNativeWrite(leaf, captureDataWritingCommand(dir.getAbsolutePath)) + val reverted = RevertNativeForTransitionHeavyStages(spark).revertToSpark(write) + val writeFiles = assertRestoredParquetWrite(reverted) + assert(writeFiles.child eq leaf, s"expected the original leaf child:\n$reverted") + } + } + + test("revertToSpark preserves a native parquet write without duplicating its unary child") { + val sparkPlan = createSparkPlan("SELECT id FROM VALUES (1), (2) AS t(id)") + val leaf = sparkPlan.collectFirst { case node: LeafExecNode => node }.getOrElse { + fail(s"expected a leaf node in test plan:\n$sparkPlan") + } + withTempPath { dir => + val write = + cometNativeWrite(cometFilter(leaf), captureDataWritingCommand(dir.getAbsolutePath)) + val reverted = RevertNativeForTransitionHeavyStages(spark).revertToSpark(write) + assertRestoredParquetWrite(reverted) + assert( + reverted.collect { case _: FilterExec => true }.size == 1, + s"expected exactly one Spark FilterExec:\n$reverted") + assert(countCometExecs(reverted) == 0, s"expected no Comet operators:\n$reverted") + } + } + + test("revertToSpark restores a native parquet write whose command has no WriteFilesExec") { + val sparkPlan = createSparkPlan("SELECT id FROM VALUES (1) AS t(id)") + val leaf = sparkPlan.collectFirst { case node: LeafExecNode => node }.getOrElse { + fail(s"expected a leaf node in test plan:\n$sparkPlan") + } + withTempPath { dir => + val command = captureDataWritingCommand(dir.getAbsolutePath) + val input = command.child match { + case writeFiles: WriteFilesExec => writeFiles.child + case other => other + } + val commandWithoutWriteFiles = + command.withNewChildren(Seq(input)).asInstanceOf[DataWritingCommandExec] + val write = cometNativeWrite(leaf, commandWithoutWriteFiles) + val reverted = RevertNativeForTransitionHeavyStages(spark).revertToSpark(write) + val restored = reverted match { + case node: DataWritingCommandExec => node + case other => fail(s"expected DataWritingCommandExec, got:\n$other") + } + assert(restored.child eq leaf, s"expected the original leaf child:\n$reverted") + assert( + restored.collect { case _: WriteFilesExec => true }.isEmpty, + s"WriteFilesExec should not be reinserted when the original command lacked it:\n$reverted") + assert( + restored.collect { case _: CometNativeWriteExec => true }.isEmpty, + s"native parquet write should be restored, not erased:\n$reverted") + } + } + + test("invalid original-plan alias skips the entire stage reversion") { + withSQLConf( + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") { + val sparkPlan = createSparkPlan("SELECT id FROM VALUES (1) AS t(id)") + val leaf = sparkPlan.collectFirst { case node: LeafExecNode => node }.getOrElse { + fail(s"expected a leaf node in test plan:\n$sparkPlan") + } + val aliasing = AliasingFallbackCometExec(leaf, leaf) + val stagePlan = CometNativeColumnarToRowExec(aliasing) + val rule = RevertNativeForTransitionHeavyStages(spark) + assert(rule.countTransitions(stagePlan) == 1) + + val result = rule(stagePlan) + + assert( + result eq stagePlan, + s"invalid fallback must leave the whole stage unchanged:\n$result") + } + } + + test("invalid original-plan alias to a Comet child skips the entire stage reversion") { + withSQLConf( + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") { + val sparkPlan = createSparkPlan("SELECT id FROM VALUES (1) AS t(id)") + val leaf = sparkPlan.collectFirst { case node: LeafExecNode => node }.getOrElse { + fail(s"expected a leaf node in test plan:\n$sparkPlan") + } + val cometChild = cometFilter(leaf) + val aliasing = AliasingFallbackCometExec(cometChild, cometChild) + val stagePlan = CometNativeColumnarToRowExec(aliasing) + val rule = RevertNativeForTransitionHeavyStages(spark) + assert(rule.countTransitions(stagePlan) == 1) + + val result = rule(stagePlan) + + assert( + result eq stagePlan, + s"invalid nested fallback must leave the whole stage unchanged:\n$result") + } + } + + test("local TopK sparkFallback returns the supplied restored child") { + val child = spark.range(20).queryExecution.sparkPlan + val originalPlan = TakeOrderedAndProjectExec(5, Seq.empty, child.output, child) + val local = CometLocalTopKExec( + Operator.newBuilder().build(), + originalPlan, + child.output, + 5, + Seq.empty, + dynamicFilterEnabled = false, + child, + SerializedPlan(None)) + val restoredChild = spark.range(10).queryExecution.sparkPlan + + assert(local.sparkFallback(Seq(restoredChild)) eq restoredChild) + intercept[CometExec.InvalidSparkFallbackException] { + local.sparkFallback(Seq.empty) + } + } + for (adaptive <- Seq(false, true)) { test(s"transition reversion preserves local TopK: AQE=$adaptive") { withSQLConf( @@ -393,6 +706,206 @@ class RevertNativeForTransitionHeavyStagesSuite extends CometTestBase { } } + test("revertToSpark leaves transitions below a shuffle that a stripped transition sat on") { + withSQLConf("spark.sql.adaptive.enabled" -> "false") { + withParquetTable((0 until 100).map(i => (i, i % 10)), "tbl") { + val df = sql("SELECT _2, count(*) FROM tbl GROUP BY _2") + df.collect() + val cometPlan = stripAQEPlan(df.queryExecution.executedPlan) + val shuffle = cometPlan + .collectFirst { case s: CometShuffleExchangeExec => s } + .getOrElse(fail(s"test requires a native shuffle:\n$cometPlan")) + // The transition that must survive lives in the map stage, under the exchange. + val preserved = + if (shuffle.child.supportsColumnar) CometNativeColumnarToRowExec(shuffle.child) + else CometSparkToColumnarExec(shuffle.child) + val shuffleWithTransition = shuffle.withNewChildren(Seq(preserved)) + // Stacked transitions sit directly on the shuffle, the shape #6152 strips through. + val onExchange = + CometSparkToColumnarExec(CometNativeColumnarToRowExec(shuffleWithTransition)) + val write = cometIcebergWrite(onExchange) + + val reverted = RevertNativeForTransitionHeavyStages(spark).revertToSpark(write) + val restoredWrite = reverted match { + case node: IcebergWriteExec => node + case other => fail(s"expected IcebergWriteExec, got:\n$other") + } + val restoredShuffle = restoredWrite.child match { + case ColumnarToRowExec(exchange: CometShuffleExchangeExec) => exchange + case exchange: CometShuffleExchangeExec => exchange + case other => + fail(s"expected the shuffle under the restored write, got:\n$other") + } + assert( + restoredShuffle.child eq preserved, + "unwrapping the transition on the shuffle must not strip the stage below it:\n" + + reverted.treeString) + } + } + } + + test("transition-heavy reversion rejects a stripped stage-boundary root") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withParquetTable((0 until 100).map(i => (i, i % 10)), "tbl") { + var cometPlan: SparkPlan = null + withSQLConf(CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false") { + val df = sql("SELECT _1, _2 FROM tbl DISTRIBUTE BY _2") + df.collect() + cometPlan = stripAQEPlan(df.queryExecution.executedPlan) + } + val exchange = cometPlan + .collectFirst { case node: CometShuffleExchangeExec => node } + .getOrElse(fail(s"test requires a native shuffle:\n$cometPlan")) + val stagePlan = CometColumnarToRowExec(exchange) + val rule = RevertNativeForTransitionHeavyStages(spark) + assert(rule.countTransitions(stagePlan) == 1) + + var result: SparkPlan = null + withSQLConf( + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") { + result = rule(stagePlan) + } + + assert( + result eq stagePlan, + "a transition-heavy stage whose stripped root is an exchange must stay unchanged:\n" + + result.treeString) + } + } + } + + test("transition-heavy revert restores row output for a columnar Spark scan") { + withSQLConf( + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0", + CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "true") { + withParquetTable((0 until 100).map(i => (i, i % 10)), "tbl") { + val (_, plan) = checkSparkAnswer("SELECT _1, _2 FROM tbl") + val executedPlan = stripAQEPlan(plan) + val resultStageRoot = unwrapCodegen(executedPlan) + + val scan = resultStageRoot match { + case transition: ColumnarToRowExec => + unwrapCodegen(transition.child) match { + case child: FileSourceScanExec => child + case other => + fail(s"expected a vectorized Spark scan under ColumnarToRow, got:\n$other") + } + case other => + fail( + "expected ColumnarToRow over a vectorized Spark scan, got " + + s"${other.getClass.getName}:\n$other") + } + assert(scan.supportsColumnar, s"the reverted scan must use its columnar path:\n$scan") + assert(countCometExecs(executedPlan) == 0, s"the stage must be reverted:\n$executedPlan") + } + } + } + + for (adaptive <- Seq(false, true); columnarRoot <- Seq(false, true)) { + test( + s"transition-heavy map-stage fallback supplies Arrow: AQE=$adaptive, columnar=$columnarRoot") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false") { + val rows = (0 until 100).map(i => (i, i % 10)) + withParquetTable(rows, "tbl") { + val df = sql("SELECT _1, _2 FROM tbl DISTRIBUTE BY _2") + df.collect() + val exchange = stripAQEPlan(df.queryExecution.executedPlan) + .collectFirst { case node: CometShuffleExchangeExec => node } + .getOrElse(fail("test requires a native shuffle")) + assert(exchange.child.isInstanceOf[CometNativeScanExec]) + val input = if (columnarRoot) exchange.child else cometFilter(exchange.child) + val stage = exchange.withNewChildren( + Seq(CometSparkToColumnarExec(CometNativeColumnarToRowExec(input)))) + val rule = RevertNativeForTransitionHeavyStages(spark) + assert(rule.countTransitions(stage.children.head) == 1) + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") { + val reverted = rule(stage).asInstanceOf[CometShuffleExchangeExec] + val bridge = reverted.child match { + case node: CometSparkToColumnarExec => node + case other => fail(s"expected an Arrow bridge after map-stage fallback:\n$other") + } + assert(bridge.child.supportsColumnar == columnarRoot) + assert(bridge.child.collect { case _: FileSourceScanExec => true }.nonEmpty) + assert(countCometExecs(bridge.child) == 0) + // Execute the native shuffle, not just its Spark fallback child: Spark batches + // satisfy supportsColumnar but cannot be cast to CometVector by the Arrow stream. + SQLExecution.withNewExecutionId(df.queryExecution) { + val actual = ColumnarToRowExec(reverted) + .executeCollect() + .map(row => (row.getInt(0), row.getInt(1))) + .toSeq + assert(actual.sorted == rows.sorted) + } + } + } + } + } + } + + for (adaptive <- Seq(false, true)) { + test(s"transition-heavy revert preserves native exchange for DISTRIBUTE BY: AQE=$adaptive") { + withSQLConf( + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.SHUFFLE_PARTITIONS.key -> "8", + "spark.sql.adaptive.coalescePartitions.enabled" -> "true", + "spark.sql.adaptive.coalescePartitions.parallelismFirst" -> "false", + "spark.sql.adaptive.advisoryPartitionSizeInBytes" -> "67108864") { + withParquetTable((0 until 100).map(i => (i, i % 10)), "tbl") { + val query = "SELECT _1, _2 FROM tbl DISTRIBUTE BY _2" + var sparkAnswer: Seq[Row] = Seq.empty + withSQLConf( + CometConf.COMET_ENABLED.key -> "false", + CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false") { + sparkAnswer = sql(query).collect().toSeq + } + val df = sql(query) + checkCometAnswer(df, sparkAnswer) + val executedPlan = stripAQEPlan(df.queryExecution.executedPlan) + val resultStageRoot = unwrapCodegen(executedPlan) + + assert( + resultStageRoot.isInstanceOf[ColumnarToRowTransition], + s"the result stage must end with a columnar-to-row transition:\n$executedPlan") + + if (adaptive) { + val read = executedPlan + .collectFirst { case node: AQEShuffleReadExec => node } + .getOrElse(fail(s"expected a coalesced AQE shuffle read:\n$executedPlan")) + assert(read.partitionSpecs.size < 8, s"shuffle must be coalesced:\n$read") + val stage = read.child match { + case node: ShuffleQueryStageExec => node + case other => fail(s"expected a shuffle query stage:\n$other") + } + assert(stage.plan.isInstanceOf[CometShuffleExchangeExec]) + } else { + val exchange = executedPlan + .collectFirst { case node: CometShuffleExchangeExec => node } + .getOrElse(fail(s"expected a native shuffle:\n$executedPlan")) + assert( + exchange.collect { case scan: CometNativeScanExec => scan }.nonEmpty, + s"the map stage must retain its native scan:\n$executedPlan") + assert( + exchange.collect { case scan: FileSourceScanExec => scan }.isEmpty, + s"fallback must not replace the scan below the exchange:\n$executedPlan") + } + } + } + } + } + test("non-AQE apply must not produce an invalid plan when the result stage reverts") { withSQLConf( CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true", diff --git a/spark/src/test/scala/org/apache/comet/serde/operator/CometIcebergNativeScanSuite.scala b/spark/src/test/scala/org/apache/comet/serde/operator/CometIcebergNativeScanSuite.scala index a11fc34d9a2..6ff5daf27fe 100644 --- a/spark/src/test/scala/org/apache/comet/serde/operator/CometIcebergNativeScanSuite.scala +++ b/spark/src/test/scala/org/apache/comet/serde/operator/CometIcebergNativeScanSuite.scala @@ -19,19 +19,19 @@ package org.apache.comet.serde.operator +import java.lang.reflect.InvocationTargetException + import org.scalatest.funsuite.AnyFunSuite import org.scalatest.matchers.should.Matchers +import org.apache.iceberg.{DeleteFile, FileContent} import org.apache.iceberg.expressions.Expressions import org.apache.spark.sql.catalyst.expressions.AttributeReference import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType, StringType, StructType} -/** - * Unit tests for [[CometIcebergNativeScan.hadoopToIcebergS3Properties]]. The pinned iceberg-rust - * S3 parser reads ONLY global `s3.*` keys (never `s3.bucket.*`), so the function drops per-bucket - * keys and promotes just the TARGET bucket's keys to global `s3.*`. Pure-function assertions, so - * a lightweight `AnyFunSuite` (no Spark session) suffices. - */ +import org.apache.comet.iceberg.IcebergReflection + +/** Unit tests for Iceberg native-scan serde helpers that do not require a Spark session. */ class CometIcebergNativeScanSuite extends AnyFunSuite with Matchers { test("complex type null residuals are not serialized") { @@ -239,4 +239,159 @@ class CometIcebergNativeScanSuite extends AnyFunSuite with Matchers { out("s3.endpoint") shouldBe "https://global.example.com" out.values.toSet should not contain "https://some.example.com" } + + private def keyMetadataMethod(clazz: Class[_]) = clazz.getMethod("keyMetadata") + + private def serializeDeleteFile(file: AnyRef) = + CometIcebergNativeScan.serializeDeleteFile( + file, + file.getClass, + file.getClass, + keyMetadataMethod(file.getClass), + _ => 0) + + test("required Iceberg delete-file accessors are present") { + Seq("content", "specId", "equalityFieldIds").foreach { accessor => + IcebergReflection.findMethod(classOf[DeleteFile], accessor).isDefined shouldBe true + } + } + + test("Iceberg delete content names match native serde literals") { + FileContent.POSITION_DELETES.toString shouldBe IcebergReflection.ContentTypes.POSITION_DELETES + FileContent.EQUALITY_DELETES.toString shouldBe IcebergReflection.ContentTypes.EQUALITY_DELETES + } + + test("position-delete file with null equality ids serializes without equality ids") { + val proto = serializeDeleteFile(new PositionDeleteFile) + proto.getContentType shouldBe "POSITION_DELETES" + proto.getPartitionSpecId shouldBe 7 + proto.getEqualityIdsCount shouldBe 0 + proto.getFilePathIdx shouldBe 0 + } + + test("equality-delete file serializes declared equality ids") { + val proto = serializeDeleteFile(new EqualityDeleteFile) + proto.getContentType shouldBe "EQUALITY_DELETES" + proto.getEqualityIdsCount shouldBe 2 + proto.getEqualityIds(0) shouldBe 3 + proto.getEqualityIds(1) shouldBe 5 + } + + test("equality-delete file with null equality ids is fatal") { + val ex = + intercept[IllegalStateException](serializeDeleteFile(new EqualityDeleteFileWithNullIds)) + ex.getMessage shouldBe + "Iceberg equality delete file 's3://bucket/eq-null-ids.parquet' has no equality field IDs" + } + + test("equality-delete file with empty equality ids is fatal") { + val ex = + intercept[IllegalStateException](serializeDeleteFile(new EqualityDeleteFileWithEmptyIds)) + ex.getMessage shouldBe + "Iceberg equality delete file 's3://bucket/eq-empty-ids.parquet' has no equality field IDs" + } + + test("content invocation failure propagates instead of defaulting to position deletes") { + val ex = + intercept[InvocationTargetException](serializeDeleteFile(new ThrowingContentDeleteFile)) + ex.getCause.getMessage shouldBe "content boom" + } + + test("spec id invocation failure propagates instead of defaulting to zero") { + val ex = + intercept[InvocationTargetException](serializeDeleteFile(new ThrowingSpecIdDeleteFile)) + ex.getCause.getMessage shouldBe "spec boom" + } + + test("equality-id invocation failure propagates instead of dropping ids") { + val ex = + intercept[InvocationTargetException](serializeDeleteFile(new ThrowingEqualityIdsDeleteFile)) + ex.getCause.getMessage shouldBe "ids boom" + } + + test("missing content accessor is fatal") { + assertThrows[NoSuchMethodException](serializeDeleteFile(new NoContentAccessorDeleteFile)) + } + + test("missing equality-id accessor is fatal") { + assertThrows[NoSuchMethodException](serializeDeleteFile(new NoEqualityIdsAccessorDeleteFile)) + } + + test("missing delete-file path accessor is fatal") { + val ex = intercept[RuntimeException](serializeDeleteFile(new NoPathAccessorDeleteFile)) + ex.getMessage should include("Neither location() nor path() is declared") + } + + private abstract class BaseDeleteFile { + def location(): String + def format(): String = "PARQUET" + def recordCount(): java.lang.Long = java.lang.Long.valueOf(1L) + def keyMetadata(): java.nio.ByteBuffer = null + } + + private class PositionDeleteFile extends BaseDeleteFile { + override def location(): String = "s3://bucket/pos-delete.parquet" + def content(): String = "POSITION_DELETES" + def specId(): Int = 7 + def equalityFieldIds(): java.util.List[Integer] = null + } + + private class EqualityDeleteFile extends BaseDeleteFile { + override def location(): String = "s3://bucket/eq-delete.parquet" + def content(): String = "EQUALITY_DELETES" + def specId(): Int = 0 + def equalityFieldIds(): java.util.List[Integer] = + java.util.List.of(Integer.valueOf(3), Integer.valueOf(5)) + } + + private class EqualityDeleteFileWithNullIds extends BaseDeleteFile { + override def location(): String = "s3://bucket/eq-null-ids.parquet" + def content(): String = "EQUALITY_DELETES" + def specId(): Int = 0 + def equalityFieldIds(): java.util.List[Integer] = null + } + + private class EqualityDeleteFileWithEmptyIds extends BaseDeleteFile { + override def location(): String = "s3://bucket/eq-empty-ids.parquet" + def content(): String = "EQUALITY_DELETES" + def specId(): Int = 0 + def equalityFieldIds(): java.util.List[Integer] = java.util.List.of[Integer]() + } + + private class ThrowingContentDeleteFile extends BaseDeleteFile { + override def location(): String = "s3://bucket/d.parquet" + def content(): String = throw new RuntimeException("content boom") + def specId(): Int = 0 + def equalityFieldIds(): java.util.List[Integer] = null + } + + private class ThrowingSpecIdDeleteFile extends BaseDeleteFile { + override def location(): String = "s3://bucket/d.parquet" + def content(): String = "POSITION_DELETES" + def specId(): Int = throw new RuntimeException("spec boom") + def equalityFieldIds(): java.util.List[Integer] = null + } + + private class ThrowingEqualityIdsDeleteFile extends BaseDeleteFile { + override def location(): String = "s3://bucket/d.parquet" + def content(): String = "EQUALITY_DELETES" + def specId(): Int = 0 + def equalityFieldIds(): java.util.List[Integer] = throw new RuntimeException("ids boom") + } + + private class NoContentAccessorDeleteFile extends BaseDeleteFile { + override def location(): String = "s3://bucket/d.parquet" + def specId(): Int = 0 + def equalityFieldIds(): java.util.List[Integer] = null + } + + private class NoEqualityIdsAccessorDeleteFile extends BaseDeleteFile { + override def location(): String = "s3://bucket/d.parquet" + def content(): String = "EQUALITY_DELETES" + def specId(): Int = 0 + } + + private class NoPathAccessorDeleteFile { + def keyMetadata(): java.nio.ByteBuffer = null + } } diff --git a/spark/src/test/scala/org/apache/comet/udf/CometUdfBridgeSuite.scala b/spark/src/test/scala/org/apache/comet/udf/CometUdfBridgeSuite.scala index 7efc5f12cbc..71ec1f856e5 100644 --- a/spark/src/test/scala/org/apache/comet/udf/CometUdfBridgeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/udf/CometUdfBridgeSuite.scala @@ -102,6 +102,8 @@ class CometUdfBridgeSuite extends AnyFunSuite { outArrays(0).memoryAddress(), outSchemas(0).memoryAddress(), numRows, + 0, + 0L, null, null) @@ -172,6 +174,8 @@ class CometUdfBridgeSuite extends AnyFunSuite { outArrays(0).memoryAddress(), outSchemas(0).memoryAddress(), numRows, + 0, + 0L, null, null) diff --git a/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala b/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala index d6237f76b8e..0485fb61818 100644 --- a/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala @@ -29,6 +29,7 @@ import scala.reflect.ClassTag import org.apache.arrow.c.ArrowArrayStream import org.apache.arrow.memory.RootAllocator +import org.apache.arrow.vector.BigIntVector import org.apache.arrow.vector.ipc.ArrowReader import org.apache.arrow.vector.types.pojo.{ArrowType, Field, Schema} import org.apache.spark.executor.TaskMetrics @@ -40,11 +41,13 @@ import org.apache.spark.sql.comet.{CometExec, CometExecUtils, CometMetricNode, C import org.apache.spark.sql.comet.execution.arrow.CometArrowStream import org.apache.spark.sql.functions.{col, udf} import org.apache.spark.sql.types.{LongType, StructField, StructType} +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} import org.apache.spark.util.JsonProtocol -import org.apache.comet.{CometConf, CometExecIterator, CometExecutorMemoryUsage, CometNativeException, CometShuffleBlockIterator, Native} +import org.apache.comet.{CometArrowAllocator, CometConf, CometExecIterator, CometExecutorMemoryUsage, CometNativeException, CometShuffleBlockIterator, Native} import org.apache.comet.serde.Config.ConfigMap import org.apache.comet.serde.OperatorOuterClass +import org.apache.comet.vector.CometPlainVector /** * Regression tests for the native plan lifecycle: every `createPlan` must be balanced by exactly @@ -237,6 +240,56 @@ class CometExecIteratorLifecycleSuite extends CometTestBase { } } + test("an input that throws after its first batch fails the plan with its own exception") { + val allocator = CometArrowAllocator.newChildAllocator("input-failure", 0, Long.MaxValue) + try { + withTaskContext(4600000L) { + val vector = new BigIntVector("test", allocator) + vector.allocateNew(2) + vector.set(0, 1L) + vector.set(1, 2L) + vector.setValueCount(2) + val first = new ColumnarBatch(Array[ColumnVector](new CometPlainVector(vector)), 2) + val boom = new IllegalStateException("injected input failure") + // The stream's schema is derived from the first batch on the JVM, so a failure there + // propagates unchanged. Native pulls every later batch through the exported stream, and + // Arrow Java hands native only the text of what that pull threw. + val input = Iterator.single(first) ++ new Iterator[ColumnarBatch] { + override def hasNext: Boolean = throw boom + override def next(): ColumnarBatch = throw boom + } + val schema = StructType(Seq(StructField("test", LongType, nullable = false))) + val stream = CometArrowStream.fromColumnarBatchIter( + input, + schema, + CometArrowStream.NATIVE_TIMEZONE, + "input-failure") + val limitOp = + CometExecUtils.getLimitNativePlan(Seq(PrettyAttribute("test", LongType)), 100).get + val iter = CometExec.getCometIterator( + Array(stream.asInstanceOf[Object]), + 1, + limitOp, + CometMetricNode(Map.empty), + 1, + 0, + None, + Seq.empty) + + assert(iter.hasNext && iter.next().numRows() == 2) + val thrown = intercept[Throwable](iter.hasNext) + assert(thrown eq boom, s"expected the input's own exception, got: $thrown") + + TaskContext.get().asInstanceOf[TaskContextImpl].markTaskCompleted(None) + assert( + CometArrowStream.inputFailure(stream).isEmpty, + "a completed task must not keep its input streams' failures") + } + } finally { + allocator.close() + } + } + test("a plan fed only by native scans returns its memory when its consumer stops early") { // See issue #2453. The sort keeps its sorted runs in memory, and the merge that reads them // spawns a Tokio task for each run. diff --git a/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala index c9b6d3de7fc..efd899252ac 100644 --- a/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala @@ -20,11 +20,13 @@ package org.apache.spark import java.util.Properties +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.{AtomicInteger, AtomicLong, AtomicReference} import org.apache.logging.log4j.{Level, LogManager} import org.apache.logging.log4j.core.{LogEvent, LoggerContext} import org.apache.spark.executor.TaskMetrics -import org.apache.spark.memory.{MemoryConsumer, TaskMemoryManager, TestMemoryManager} +import org.apache.spark.memory.{MemoryConsumer, MemoryMode, TaskMemoryManager, TestMemoryManager, UnifiedMemoryManager} class CometTaskMemoryManagerSuite extends SparkFunSuite { @@ -152,10 +154,166 @@ class CometTaskMemoryManagerSuite extends SparkFunSuite { appender.loggingEvents.toSeq } + test("an acquire waiting in Spark survives a release that empties the task's balance") { + val memoryManager = offHeapMemoryManager() + val otherTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 1L)) + + withTaskContext(new TaskMemoryManager(memoryManager, 0L)) { taskMemoryManager => + val manager = new CometTaskMemoryManager(1L, 0L) + assert(otherTask.acquireMemory(90L) == 90L) + assert(manager.acquireMemory(10L) == 10L) + + // Another native thread releases the task's last 10 bytes while the acquire waits. + acquireWhileBalanceEmpties(manager, taskMemoryManager, otherTask) { + manager.releaseMemory(10L) + } + } + } + + test("an acquire waiting in Spark survives a sibling consumer freeing the task's last bytes") { + val memoryManager = offHeapMemoryManager() + val otherTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 1L)) + + withTaskContext(new TaskMemoryManager(memoryManager, 0L)) { taskMemoryManager => + val manager = new CometTaskMemoryManager(1L, 0L) + val sibling = new OffHeapConsumer(taskMemoryManager) + assert(otherTask.acquireMemory(90L) == 90L) + assert(sibling.acquireMemory(10L) == 10L) + + // A JVM consumer of the same task frees with freeMemory, which takes no task monitor. + acquireWhileBalanceEmpties(manager, taskMemoryManager, otherTask) { + sibling.freeMemory(10L) + } + } + } + + test("a NoSuchElementException other than Spark's missing task entry is rethrown") { + val wrongTask = missingEntryError("key not found: 7", fromExecutionMemoryPool = true) + val notFromPool = missingEntryError("key not found: 0", fromExecutionMemoryPool = false) + for (error <- Seq(wrongTask, notFromPool)) { + val taskMemoryManager = new FailingTaskMemoryManager(error) + withTaskContext(taskMemoryManager) { _ => + val manager = new CometTaskMemoryManager(1L, 0L) + val thrown = intercept[NoSuchElementException](manager.acquireMemory(10L)) + assert(thrown eq error) + assert(taskMemoryManager.calls.get == 1) + assert(manager.getUsed == 0L) + } + } + } + + test("an acquire that keeps losing its task entry gives up with a zero grant") { + val error = missingEntryError("key not found: 0", fromExecutionMemoryPool = true) + val taskMemoryManager = new FailingTaskMemoryManager(error) + withTaskContext(taskMemoryManager) { _ => + val manager = new CometTaskMemoryManager(1L, 0L) + assert(manager.acquireMemory(10L) == 0L) + assert(taskMemoryManager.calls.get == maxAcquireAttempts) + assert(manager.getUsed == 0L) + assert(nativeMemoryConsumer(manager).getUsed == 0L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + } + } + + private val TimeoutSeconds = 10L + + /** + * The task holds 10 bytes of a 100 byte pool and another task holds 90. An acquire of 20 bytes + * waits in Spark below its minimum share of 25 while `release` takes the task's balance to + * zero, which removes the task's entry from Spark's pool. Once the other task frees its memory + * the acquire must be granted all 20 bytes. + */ + private def acquireWhileBalanceEmpties( + manager: CometTaskMemoryManager, + taskMemoryManager: TaskMemoryManager, + otherTask: OffHeapConsumer)(release: => Unit): Unit = { + val granted = new AtomicLong(-1L) + val failure = new AtomicReference[Throwable]() + val acquire = new Thread(() => + try granted.set(manager.acquireMemory(20L)) + catch { case t: Throwable => failure.set(t) }) + acquire.setDaemon(true) + + try { + acquire.start() + val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(TimeoutSeconds) + while (!waitingInSpark(acquire) && System.nanoTime() < deadline) Thread.sleep(10) + assert(waitingInSpark(acquire), s"the acquire is ${acquire.getState}") + release + } finally { + // Free the other task's memory so that the acquire does not outlive a failed test. + otherTask.freeMemory(otherTask.getUsed) + acquire.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds)) + } + + assert(!acquire.isAlive, s"the acquire is ${acquire.getState}") + assert(failure.get == null, s"the acquire failed: ${failure.get}") + assert(granted.get == 20L) + assert(manager.getUsed == 20L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 20L) + manager.releaseMemory(20L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + } + + private def waitingInSpark(thread: Thread): Boolean = + thread.getState == Thread.State.WAITING && + thread.getStackTrace.exists(_.getClassName == "org.apache.spark.memory.ExecutionMemoryPool") + + /** A 100 byte off-heap execution pool. */ + private def offHeapMemoryManager(): UnifiedMemoryManager = { + val conf = new SparkConf() + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", "100") + .set("spark.memory.storageFraction", "0") + new UnifiedMemoryManager(conf, 1000L, 500L, 1) + } + + /** An off-heap consumer that never spills. */ + private class OffHeapConsumer(taskMemoryManager: TaskMemoryManager) + extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP) { + override def spill(size: Long, trigger: MemoryConsumer): Long = 0L + } + + private val MaxFailingCalls = 10 + + /** A task memory manager whose every acquire throws `error`. */ + private class FailingTaskMemoryManager(error: Throwable) + extends TaskMemoryManager(new TestMemoryManager(new SparkConf()), 0L) { + val calls = new AtomicInteger() + + override def acquireExecutionMemory(required: Long, consumer: MemoryConsumer): Long = { + // Past this bound a caller is retrying without a cap; fail the test instead of hanging it. + if (calls.incrementAndGet() > MaxFailingCalls) { + throw new IllegalStateException( + s"acquireExecutionMemory called more than $MaxFailingCalls times") + } + throw error + } + } + + private def missingEntryError( + message: String, + fromExecutionMemoryPool: Boolean): NoSuchElementException = { + val error = new NoSuchElementException(message) + if (fromExecutionMemoryPool) { + val poolFrame = new StackTraceElement( + "org.apache.spark.memory.ExecutionMemoryPool", + "acquireMemory", + "ExecutionMemoryPool.scala", + 115) + error.setStackTrace(poolFrame +: error.getStackTrace) + } + error + } + private def withTaskMemoryManager(f: TaskMemoryManager => Unit): Unit = { val memoryManager = new TestMemoryManager(new SparkConf()) memoryManager.limit(1024) - val taskMemoryManager = new TaskMemoryManager(memoryManager, 0L) + withTaskContext(new TaskMemoryManager(memoryManager, 0L))(f) + } + + private def withTaskContext(taskMemoryManager: TaskMemoryManager)( + f: TaskMemoryManager => Unit): Unit = { val taskContext = new TaskContextImpl( stageId = 0, stageAttemptNumber = 0, @@ -182,6 +340,12 @@ class CometTaskMemoryManagerSuite extends SparkFunSuite { } } + private def maxAcquireAttempts: Int = { + val field = classOf[CometTaskMemoryManager].getDeclaredField("MAX_ACQUIRE_ATTEMPTS") + field.setAccessible(true) + field.getInt(null) + } + private def nativeMemoryConsumer(manager: CometTaskMemoryManager): MemoryConsumer = { val field = classOf[CometTaskMemoryManager].getDeclaredField("nativeMemoryConsumer") field.setAccessible(true) diff --git a/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala b/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala new file mode 100644 index 00000000000..b978c272802 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala @@ -0,0 +1,285 @@ +/* + * 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.shuffle.comet + +import java.util.Properties +import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.{AtomicInteger, AtomicReference} + +import scala.collection.mutable.ArrayBuffer + +import org.apache.logging.log4j.{Level, LogManager} +import org.apache.logging.log4j.core.{LogEvent, LoggerContext} +import org.apache.spark.{CometTaskMemoryManager, SparkConf, SparkFunSuite, TaskContext, TaskContextImpl} +import org.apache.spark.executor.TaskMetrics +import org.apache.spark.memory.{MemoryConsumer, MemoryMode, SparkOutOfMemoryError, TaskMemoryManager, UnifiedMemoryManager} + +class CometUnifiedShuffleMemoryAllocatorSuite extends SparkFunSuite { + + private val PageSize = 20L + private val TimeoutSeconds = 10L + private val MaxAllocateAttempts = 3 + + test("a page waiting in Spark survives a release that empties the task's balance") { + checkWaitingAllocationSurvives(20L)(allocatePage) + } + + test("a pointer array waiting in Spark survives a release that empties the task's balance") { + checkWaitingAllocationSurvives(24L)(allocatePointerArray) + } + + test("a NoSuchElementException other than Spark's missing task entry is rethrown") { + val otherMessage = noSuchElement("no entry for task 0", fromExecutionMemoryPool = true) + val notFromPool = noSuchElement("key not found: 0", fromExecutionMemoryPool = false) + for (error <- Seq(otherMessage, notFromPool); allocation <- allocations) { + val taskMemoryManager = new FailingTaskMemoryManager(() => error, failures = Int.MaxValue) + val allocator = new CometUnifiedShuffleMemoryAllocator(taskMemoryManager, PageSize) + val events = logEvents(Level.INFO) { + val thrown = intercept[NoSuchElementException](allocation(allocator)) + assert(thrown eq error) + } + assert(events.isEmpty, messages(events)) + assert(taskMemoryManager.calls.get == 1) + assert(allocator.getUsed == 0L) + assert(taskMemoryManager.cleanUpAllAllocatedMemory() == 0L) + } + } + + test("an allocation that keeps losing its task entry fails with Spark's out of memory error") { + for ((allocation, requested) <- allocations.zip(Seq("20", "24"))) { + // A new exception on every call, so the test can tell which one the allocator keeps. + val taskMemoryManager = new FailingTaskMemoryManager( + () => noSuchElement("key not found: 0", fromExecutionMemoryPool = true), + failures = Int.MaxValue) + val allocator = new CometUnifiedShuffleMemoryAllocator(taskMemoryManager, PageSize) + val events = logEvents(Level.INFO) { + val thrown = intercept[SparkOutOfMemoryError](allocation(allocator)) + assert(thrown.getErrorClass == "UNABLE_TO_ACQUIRE_MEMORY") + assert(thrown.getMessageParameters.get("requestedBytes") == requested) + assert(thrown.getMessageParameters.get("receivedBytes") == "0") + // The last attempt's exception carries the task id and where Spark was waiting. + assert(thrown.getCause eq taskMemoryManager.errors.last) + } + assert(events.map(_.getLevel) == Seq(Level.INFO, Level.INFO, Level.WARN), messages(events)) + assert(events.last.getThrown eq taskMemoryManager.errors.last) + assert(taskMemoryManager.calls.get == MaxAllocateAttempts) + assert(taskMemoryManager.errors.distinct.size == MaxAllocateAttempts) + assert(allocator.getUsed == 0L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + assert(taskMemoryManager.cleanUpAllAllocatedMemory() == 0L) + } + } + + test("an allocation that loses its task entry twice is granted on the last attempt") { + val error = noSuchElement("key not found: 0", fromExecutionMemoryPool = true) + for ((allocation, requested) <- allocations.zip(Seq(20L, 24L))) { + val taskMemoryManager = + new FailingTaskMemoryManager(() => error, failures = MaxAllocateAttempts - 1) + val allocator = new CometUnifiedShuffleMemoryAllocator(taskMemoryManager, PageSize) + val events = logEvents(Level.INFO) { + assert(allocation(allocator) == requested) + } + assert(events.map(_.getLevel) == Seq(Level.INFO, Level.INFO), messages(events)) + assert(taskMemoryManager.calls.get == MaxAllocateAttempts) + assert(allocator.getUsed == 0L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + assert(taskMemoryManager.cleanUpAllAllocatedMemory() == 0L) + } + } + + /** Allocates a 20 byte page, frees it and returns its size. */ + private def allocatePage(allocator: CometUnifiedShuffleMemoryAllocator): Long = { + val page = allocator.allocate(20L) + val size = page.size() + allocator.free(page) + size + } + + /** Allocates a 3 entry (24 byte) pointer array, frees it and returns its size in bytes. */ + private def allocatePointerArray(allocator: CometUnifiedShuffleMemoryAllocator): Long = { + val array = allocator.allocateArray(3L) + val size = array.memoryBlock().size() + allocator.freeArray(array) + size + } + + private def allocations: Seq[CometUnifiedShuffleMemoryAllocator => Long] = + Seq(allocatePage _, allocatePointerArray _) + + /** + * The task holds 10 bytes of a 100 byte pool through a native consumer and another task holds + * 90. `allocate` asks for `requested` bytes on another thread and waits in Spark below the + * task's minimum share of 25 while the native consumer releases the task's last 10 bytes, which + * removes the task's entry from Spark's pool. Once the other task frees its memory the + * allocation must be granted in full. + */ + private def checkWaitingAllocationSurvives(requested: Long)( + allocate: CometUnifiedShuffleMemoryAllocator => Long): Unit = { + val memoryManager = offHeapMemoryManager() + val otherTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 1L)) + + withTaskContext(new TaskMemoryManager(memoryManager, 0L)) { taskMemoryManager => + val native = new CometTaskMemoryManager(1L, 0L) + val allocator = new CometUnifiedShuffleMemoryAllocator(taskMemoryManager, PageSize) + // The other task goes first: acquiring its 90 bytes after the task's 10 would cap it at 50. + assert(otherTask.acquireMemory(90L) == 90L) + assert(native.acquireMemory(10L) == 10L) + + val granted = new AtomicReference[java.lang.Long]() + val failure = new AtomicReference[Throwable]() + val thread = new Thread(() => + try granted.set(allocate(allocator)) + catch { case t: Throwable => failure.set(t) }) + thread.setDaemon(true) + + val events = logEvents(Level.INFO) { + try { + thread.start() + val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(TimeoutSeconds) + while (!waitingInSpark(thread) && System.nanoTime() < deadline) Thread.sleep(10) + assert(waitingInSpark(thread), s"the allocation is ${thread.getState}") + native.releaseMemory(10L) + } finally { + // Free the other task's memory so that the allocation does not outlive a failed test. + otherTask.freeMemory(otherTask.getUsed) + thread.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds)) + } + } + + assert(!thread.isAlive, s"the allocation is ${thread.getState}") + assert(failure.get == null, s"the allocation failed: ${failure.get}") + assert(granted.get == requested) + assert(allocator.getUsed == 0L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + // The first attempt lost the task's entry and the allocator tried again. + assert(events.map(_.getLevel) == Seq(Level.INFO), messages(events)) + } + } + + private def waitingInSpark(thread: Thread): Boolean = + thread.getState == Thread.State.WAITING && + thread.getStackTrace.exists(_.getClassName == "org.apache.spark.memory.ExecutionMemoryPool") + + /** A 100 byte off-heap execution pool. */ + private def offHeapMemoryManager(): UnifiedMemoryManager = { + val conf = new SparkConf() + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", "100") + .set("spark.memory.storageFraction", "0") + new UnifiedMemoryManager(conf, 1000L, 500L, 1) + } + + /** An off-heap consumer that never spills. */ + private class OffHeapConsumer(taskMemoryManager: TaskMemoryManager) + extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP) { + override def spill(size: Long, trigger: MemoryConsumer): Long = 0L + } + + private val MaxFailingCalls = 10 + + /** + * A task memory manager over a 100 byte off-heap pool whose first `failures` acquires throw the + * result of `newError`, recording each one in `errors`. + */ + private class FailingTaskMemoryManager(newError: () => Throwable, failures: Int) + extends TaskMemoryManager(offHeapMemoryManager(), 0L) { + val calls = new AtomicInteger() + val errors = new ArrayBuffer[Throwable]() + + override def acquireExecutionMemory(required: Long, consumer: MemoryConsumer): Long = { + val call = calls.incrementAndGet() + // Past this bound a caller is retrying without a cap; fail the test instead of hanging it. + if (call > MaxFailingCalls) { + throw new IllegalStateException( + s"acquireExecutionMemory called more than $MaxFailingCalls times") + } + if (call <= failures) { + val error = newError() + errors += error + throw error + } + super.acquireExecutionMemory(required, consumer) + } + } + + private def noSuchElement( + message: String, + fromExecutionMemoryPool: Boolean): NoSuchElementException = { + val error = new NoSuchElementException(message) + if (fromExecutionMemoryPool) { + val poolFrame = new StackTraceElement( + "org.apache.spark.memory.ExecutionMemoryPool", + "acquireMemory", + "ExecutionMemoryPool.scala", + 115) + error.setStackTrace(poolFrame +: error.getStackTrace) + } + error + } + + /** The events that the allocator logs at `level` or above. */ + private def logEvents(level: Level)(f: => Unit): Seq[LogEvent] = { + val appender = new LogAppender("unified shuffle memory allocator") + appender.setThreshold(level) + val loggers = Seq(classOf[CometUnifiedShuffleMemoryAllocator].getName) + val context = LogManager.getContext(false).asInstanceOf[LoggerContext] + val unconfigured = loggers.filterNot(context.getConfiguration.getLoggers.containsKey) + try withLogAppender(appender, loggers, Some(level))(f) + finally { + // For a logger with no config of its own, withLogAppender adds one and never removes it. + // It copies the root config's additivity, which is off, so left in place it would keep the + // logger's events out of the test log for the rest of the run. + unconfigured.foreach(context.getConfiguration.removeLogger) + context.updateLoggers() + } + appender.loggingEvents.toSeq + } + + private def messages(events: Seq[LogEvent]): String = + events.map(_.getMessage.getFormattedMessage).mkString("\n") + + private def withTaskContext(taskMemoryManager: TaskMemoryManager)( + f: TaskMemoryManager => Unit): Unit = { + val taskContext = new TaskContextImpl( + stageId = 0, + stageAttemptNumber = 0, + partitionId = 0, + numPartitions = 1, + taskAttemptId = 0L, + attemptNumber = 0, + taskMemoryManager = taskMemoryManager, + localProperties = new Properties, + metricsSystem = null, + taskMetrics = TaskMetrics.empty, + cpus = 1, + resources = Map.empty) + + TaskContext.setTaskContext(taskContext) + try { + f(taskMemoryManager) + } finally { + try { + taskMemoryManager.cleanUpAllAllocatedMemory() + } finally { + TaskContext.unset() + } + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala index 07864c6fec8..60b646fd74e 100644 --- a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala +++ b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala @@ -79,9 +79,9 @@ abstract class CometTestBase } /** - * Sets each Spark-to-Arrow conversion that [[sparkConf]] turns on, which are off by default, to - * `enabled`. These are the conversions `spark.comet.sparkToColumnar.enabled` stood for before - * each had a config of its own. + * Sets each Spark-to-Arrow conversion that [[sparkConf]] turns on to `enabled`. All but the + * `OneRowRelation` conversion are off by default. These are the conversions + * `spark.comet.sparkToColumnar.enabled` stood for before each had a config of its own. */ protected def sparkToArrowConversionConfs(enabled: Boolean): Seq[(String, String)] = Seq( diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExplodeBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExplodeBenchmark.scala index 86909f839c5..faafcb4d98a 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExplodeBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExplodeBenchmark.scala @@ -22,6 +22,8 @@ package org.apache.spark.sql.benchmark import java.io.File import org.apache.spark.sql.catalyst.optimizer.InferFiltersFromGenerate +import org.apache.spark.sql.catalyst.plans.logical.Generate +import org.apache.spark.sql.comet.CometExplodeExec import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf @@ -59,11 +61,8 @@ import org.apache.comet.CometConf * processes. Worth knowing when reading these numbers: real queries do get that filter, so a * plain `explode` in production usually sees an array column with no nulls and no empty rows. * - * Only array inputs are covered. Comet declines to convert a generator over a map - * (https://github.com/apache/datafusion-comet/issues/2837), so the Comet arm of such a case would - * be Spark's `GenerateExec` behind a columnar-to-row transition and its timing would say nothing - * about `CometExplodeExec`. The nesting group below therefore reaches its event list through an - * `array>` where a map would be the more natural modeling choice. + * Map cases consume both keys and values. The nesting group models event lists as a map keyed by + * platform, including a second generator over the entries in each map value. */ object CometExplodeBenchmark extends CometBenchmarkBase { @@ -109,6 +108,11 @@ object CometExplodeBenchmark extends CometBenchmarkBase { |END""".stripMargin } + /** Maps with the same NULL/empty distribution as the array cases and ordered platform keys. */ + private def mapColumn(elementExpr: String, len: Int): String = + s"map_from_arrays(${arrayColumn("concat('p_', CAST(x AS STRING))", len)}, " + + s"${arrayColumn(elementExpr, len)})" + /** * Types `expr` as nullable without ever evaluating to null, `guard` being a predicate that is * never true. @@ -153,20 +157,19 @@ object CometExplodeBenchmark extends CometBenchmarkBase { | 'page', concat('page_', CAST(id + y AS STRING)), | 'source', concat('src_', CAST(id % 4 AS STRING)))""".stripMargin - /** One platform's event list, the element of the array the nesting group explodes. */ + /** One platform's event list, the value of the map the nesting group explodes. */ private val eventStruct = s"""named_struct( - | 'platform', concat('p_', CAST(x AS STRING)), | 'entries', ${fullArray(entryStruct, entriesLen, "y")})""".stripMargin - private val eventsArray = arrayColumn(eventStruct, eventsLen) + private val eventsMap = mapColumn(eventStruct, eventsLen) - /** [[eventsArray]] wrapped in the struct chain that puts it eight accessors down. */ + /** [[eventsMap]] wrapped in the struct chain that puts it eight accessors down. */ private val profileColumn = { val sessions = s"""named_struct( | 'device', named_struct('type', 'mobile', 'os', 'iOS'), - | 'events', $eventsArray)""".stripMargin + | 'events', $eventsMap)""".stripMargin val account = accountPath.foldRight(s"named_struct('sessions', $sessions)") { (field, inner) => s"named_struct('$field', $inner)" } @@ -179,16 +182,19 @@ object CometExplodeBenchmark extends CometBenchmarkBase { /** * The temp views the benchmark reads. * - * Each array column gets its own view rather than sharing one wide table, so that a case is - * never charged for scanning an array column it does not read. The nesting group is the - * exception: `events` and `profile` hold the same array at two different depths, and keeping - * them in one view is what makes them the same array. Nested schema pruning stops either case - * from reading the other's copy. + * Each collection column gets its own view rather than sharing one wide table, so that a case + * is never charged for scanning a column it does not read. The nesting group is the exception: + * `events` and `profile` hold the same map at two different depths, and keeping them in one + * view is what makes them the same map. Nested schema pruning stops either case from reading + * the other's copy. */ private val views: Seq[TempView] = Seq( TempView("arr_len2", numRows, Seq(s"${arrayColumn("id + x", 2)} AS arr")), TempView("arr_len10", numRows, Seq(s"${arrayColumn("id + x", 10)} AS arr")), TempView("arr_len100", numRows, Seq(s"${arrayColumn("id + x", 100)} AS arr")), + TempView("map_len2", numRows, Seq(s"${mapColumn("id + x", 2)} AS m")), + TempView("map_len10", numRows, Seq(s"${mapColumn("id + x", 10)} AS m")), + TempView("map_len100", numRows, Seq(s"${mapColumn("id + x", 100)} AS m")), TempView("arr_str10", numRows, Seq(s"${arrayColumn(stringElement, 10)} AS arr")), TempView( "arr_struct10", @@ -208,7 +214,7 @@ object CometExplodeBenchmark extends CometBenchmarkBase { Seq( s"${nullableExpr("id", "id < 0")} AS k", s"${nullableExpr("concat('r_', CAST(id % 8 AS STRING))", "id < 0")} AS region", - s"$eventsArray AS events", + s"$eventsMap AS events", s"$profileColumn AS profile"))) /** Writes a view's rows to Parquet and registers it. */ @@ -219,8 +225,8 @@ object CometExplodeBenchmark extends CometBenchmarkBase { } /** - * A case that applies `generator` to `arrayExpr` over `view`, carries `carried` through the - * generator alongside it, and aggregates every column that comes out. + * A case that applies `generator` to `collectionExpr` over `view`, carries `carried` through + * the generator alongside it, and aggregates every column that comes out. * * The position column of the `posexplode` variants is summed rather than counted, because * counting it does not read a position. `pos` is declared non-null, so `NullPropagation` @@ -230,24 +236,26 @@ object CometExplodeBenchmark extends CometBenchmarkBase { * `ListPositionsExpr` and unnests it alongside the values — so a sink that ignores them is not * measuring the generator it names. `sum` over the same column is value-dependent. * - * Everything else is counted rather than summed: the element type varies across these cases and - * most of the element types cannot be summed, and unlike `pos` the generated column is nullable - * everywhere, so the count is not rewritten away. + * Map keys are also non-nullable. Sum their lengths to keep the key column in the plan and + * consume its bytes. Values and carried columns are nullable and counted, since the value types + * vary and most cannot be summed. */ private def generatorCase( generator: String, view: String, len: Int, rows: Int = numRows, - arrayExpr: String = "arr", + collectionExpr: String = "arr", + mapInput: Boolean = false, carried: Seq[String] = Nil, where: Option[String] = None): Case = { val position = generator.startsWith("posexplode") - val generated = if (position) Seq("pos", "col") else Seq("col") + val values = if (mapInput) Seq("key", "col") else Seq("col") + val generated = if (position) "pos" +: values else values val alias = if (generated.length == 1) s"AS ${generated.head}" else generated.mkString("AS (", ", ", ")") - val projectList = (carried :+ s"$generator($arrayExpr) $alias").mkString(", ") + val projectList = (carried :+ s"$generator($collectionExpr) $alias").mkString(", ") val filter = where.map(w => s" WHERE $w").getOrElse("") val aggregates = (carried ++ generated).map(aggregate) val query = @@ -262,14 +270,20 @@ object CometExplodeBenchmark extends CometBenchmarkBase { // Positions run 0 until len on every row that emits, and are null on the rows only an // outer variant emits, which `sum` skips. case "pos" => elements * (len - 1) / 2 + case "key" => nonEmptyRows(rows) * (1 to len).map(i => s"p_$i".length.toLong).sum case _ => elements } Case(query, expected) } - /** See [[generatorCase]] for why `pos` alone is summed. */ - private def aggregate(column: String): String = - if (column == "pos") s"sum($column)" else s"count($column)" + /** + * Map keys are non-nullable, so consume their bytes rather than letting count become count(1). + */ + private def aggregate(column: String): String = column match { + case "pos" => s"sum($column)" + case "key" | "platform" => s"sum(length($column))" + case _ => s"count($column)" + } /** * Runs `query` under both engines, untimed, and fails unless each produces exactly `expected`. @@ -290,7 +304,8 @@ object CometExplodeBenchmark extends CometBenchmarkBase { withSQLConf( CometConf.COMET_ENABLED.key -> cometEnabled.toString, CometConf.COMET_EXEC_ENABLED.key -> cometEnabled.toString) { - val actual = spark.sql(query).collect().head.toSeq.map { + val df = spark.sql(query) + val actual = df.collect().head.toSeq.map { case null => null case value: Number => value.longValue() case other => other @@ -300,6 +315,18 @@ object CometExplodeBenchmark extends CometBenchmarkBase { s"$name: $engine produced $actual, expected $expected. The case is not aggregating " + s"what it names.\n$query") } + if (cometEnabled) { + val expectedGenerators = df.queryExecution.optimizedPlan.collect { case _: Generate => + 1 + }.size + val nativeGenerators = collect(df.queryExecution.executedPlan) { + case e: CometExplodeExec => e + }.size + assert( + expectedGenerators > 0 && nativeGenerators == expectedGenerators, + s"$name executed $nativeGenerators of $expectedGenerators generators natively:\n" + + df.queryExecution.executedPlan) + } } } @@ -338,6 +365,24 @@ object CometExplodeBenchmark extends CometBenchmarkBase { } } + runBenchmark("Explode - maps") { + Seq(2, 10, 100).foreach { len => + runCase( + s"explode map[$len]", + generatorCase( + "explode", + s"map_len$len", + len, + collectionExpr = "m", + mapInput = true)) + } + Seq("posexplode", "explode_outer", "posexplode_outer").foreach { generator => + runCase( + s"$generator map[10]", + generatorCase(generator, "map_len10", 10, collectionExpr = "m", mapInput = true)) + } + } + runBenchmark("Explode - element type") { Seq("bigint" -> "arr_len10", "string" -> "arr_str10", "struct" -> "arr_struct10") .foreach { case (elementType, view) => @@ -382,17 +427,15 @@ object CometExplodeBenchmark extends CometBenchmarkBase { * * The shape is the one asked for in review: a customer profile whose event list sits eight * struct accessors down, holding a second array of four-field structs inside each element. The - * requested outer container was a map keyed by platform, which is where a real schema would put - * it; that is an `array>` here because Comet has no native generator - * over maps yet (#2837) and the Comet arm would silently be Spark. + * outer container is a `map>>>` keyed by platform. * * The whole event struct is counted rather than one of its fields, so nested schema pruning * cannot narrow the exploded element and leave the case measuring a two-column gather. It does * prune the siblings no case reads: on the executed plan the `depth 8` scan's `profile` column - * is `struct>>>` with `billing`, - * `addresses` and `device` gone, and the `depth 1` scan does not project `profile` at all. + * is `struct>>>` with `billing`, `addresses` + * and `device` gone, and the `depth 1` scan does not project `profile` at all. * - * `depth 1` and `depth 8` explode the same array, written twice into the same file, so the pair + * `depth 1` and `depth 8` explode the same map, written twice into the same file, so the pair * measures what the struct chain costs: extra definition levels in the Parquet column and a * chain of `GetStructField` above the scan. `then its inner array` chains a second generator * onto the first, which is the shape the review asked about; its fan-out is 20 against the @@ -401,13 +444,14 @@ object CometExplodeBenchmark extends CometBenchmarkBase { private def nestedInputCases: Seq[(String, Case)] = { val carried = Seq("k", "region") val single = Seq("depth 1" -> "events", "depth 8" -> deepEvents).map { case (label, path) => - s"explode array[$eventsLen] at $label" -> + s"explode map[$eventsLen] at $label" -> generatorCase( "explode", "nested", eventsLen, rows = nestedRows, - arrayExpr = path, + collectionExpr = path, + mapInput = true, carried = carried) } @@ -416,13 +460,13 @@ object CometExplodeBenchmark extends CometBenchmarkBase { val events = nonEmptyRows(nestedRows) * eventsLen val entries = events * entriesLen val outer = - "SELECT k, region, ev.platform AS platform, ev.entries AS entries " + - s"FROM (SELECT k, region, explode($deepEvents) AS ev FROM nested)" + "SELECT k, region, platform, ev.entries AS entries " + + s"FROM (SELECT k, region, explode($deepEvents) AS (platform, ev) FROM nested)" val chained = s"SELECT ${Seq("k", "region", "platform", "entry").map(aggregate).mkString(", ")} " + s"FROM (SELECT k, region, platform, explode(entries) AS entry FROM ($outer))" - single :+ (s"explode array[$eventsLen] at depth 8, then its inner array" -> - Case(chained, Seq.fill(4)(entries))) + single :+ (s"explode map[$eventsLen] at depth 8, then its inner array" -> + Case(chained, Seq(entries, entries, entries * 3, entries))) } } diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergWriteBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergWriteBenchmark.scala index fbcde2d46d6..7ab0d9d2161 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergWriteBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergWriteBenchmark.scala @@ -529,6 +529,7 @@ object CometIcebergWriteBenchmark extends CometBenchmarkBase { qe: QueryExecution, exception: Exception): Unit = {} } + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) spark.listenerManager.register(listener) try { action diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala index 046fe3b822c..230126a7025 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometInMemoryCacheBenchmark.scala @@ -28,7 +28,9 @@ import org.apache.spark.sql.catalyst.expressions.AttributeReference import org.apache.spark.sql.comet.CometInMemoryTableScanExec import org.apache.spark.sql.comet.execution.arrow.{ArrowCachedBatchSerializer, CometCachedBatchHelper} import org.apache.spark.sql.execution.ColumnarToRowExec -import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, DefaultCachedBatchSerializer, InMemoryRelation, InMemoryTableScanExec} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec +import org.apache.spark.sql.execution.aggregate.HashAggregateExec +import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, DefaultCachedBatch, DefaultCachedBatchSerializer, InMemoryRelation, InMemoryTableScanExec} import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.spark.sql.types.{DataType, LongType, StringType} @@ -236,7 +238,9 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { scanned = 6) runCodecBenchmark(flatRelation) + runBuildBenchmark(flatRelation) runSparkOperatorBenchmark(flatRelation) + runAdaptiveBenchmark(flatRelation) } runWideSparkOperatorBenchmark() } @@ -344,17 +348,101 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { } } + /** + * Building the cache in each format, and the footprint each format leaves, against Spark's own + * cache format. Turning the feature on changes how every cached relation is written, and a job + * that caches a relation to read it only once or twice pays more for the build than for the + * reads. + * + * What the build costs depends on what produces the rows being cached. Comet's format is + * written straight from the Arrow batches of a plan that runs natively, and from rows + * otherwise, while Spark's format is always built from rows, so above Comet operators it pays a + * columnar-to-row transition first. The relation's source stands in for both kinds of input: + * with spark.comet.exec.range.enabled Comet generates and projects its rows natively, as above + * a native scan, and without it the source runs on Spark. Comet and AQE are on and Comet's + * other settings are at their defaults, as in runAdaptiveBenchmark. + * + * Only one copy is cached at a time, as in runCodecBenchmark. + */ + private def runBuildBenchmark(relation: CachedRelation): Unit = { + val view = s"${relation.table}_build" + val formats = Seq( + ("Spark's cache format", sparkSerializer, Seq.empty[(String, String)]), + ("Comet's cache format, zstd", cometSerializer, codecConf("zstd")), + ("Comet's cache format, none", cometSerializer, codecConf("none"))) + val inputs = Seq( + ("Comet operators", true, Seq(CometConf.COMET_EXEC_RANGE_ENABLED.key -> "true")), + ("Spark operators", false, Seq.empty[(String, String)])) + + spark.catalog.clearCache() + withTempTable(view) { + spark + .sql(s"SELECT ${relation.columns.mkString(", ")} FROM ${relation.source}") + .createOrReplaceTempView(view) + + inputs.foreach { case (input, nativeInput, inputConf) => + val benchmark = + new Benchmark( + s"in-memory cache build from $input", + relation.rows.toLong, + output = output) + val footprints = new Array[Long](formats.length) + formats.zipWithIndex.foreach { case ((name, serializer, formatConf), i) => + // Timed around the caching alone: dropping the previous copy is setup, as in + // runCodecBenchmark. + benchmark.addTimerCase(name) { timer => + spark.catalog.uncacheTable(view) + withCacheSerializer(serializer) { + withSQLConf(adaptiveConf ++ inputConf ++ formatConf: _*) { + timer.startTiming() + spark.catalog.cacheTable(view) + spark.table(view).count() + timer.stopTiming() + } + } + // On the case's first call, which is a warmup iteration. + if (footprints(i) == 0L) { + verifyBuild(view, serializer, nativeInput) + footprints(i) = cachedBytes(view) + } + } + } + benchmark.run() + + // As in runCodecBenchmark, footprint has no column in a Benchmark table. + formats.zip(footprints).foreach { case ((name, _, _), bytes) => + benchmark.out.println(f"Cached footprint ($name): ${bytes / (1024.0 * 1024.0)}%.1f MiB") + } + } + + spark.catalog.uncacheTable(view) + } + } + + // Pins what a build case claims: the named serializer cached the relation, from a plan that ran + // natively exactly when the case says its input came from Comet operators. + private def verifyBuild(view: String, serializer: String, nativeInput: Boolean): Unit = { + val relation = cachedRelation(view) + val actual = relation.cacheBuilder.serializer.getClass.getName + assert(actual == serializer, s"Expected a relation cached by $serializer, not $actual") + val plan = relation.cacheBuilder.cachedPlan + val cometOperators = collect(plan) { case p if p.nodeName.startsWith("Comet") => p } + assert( + cometOperators.nonEmpty == nativeInput, + s"Expected the cached plan to run ${if (nativeInput) "natively" else "on Spark"}:\n$plan") + } + /** * Reads that feed Spark operators rather than Comet ones, against Spark's own cache format. * * Native execution is off in every case, so this measures Spark consuming the cached data: the - * shape where Comet's format has something to lose, and the reason the feature is off by - * default. Comet's format is read two ways. With Comet off, every Spark operator reads rows - * from the cache scan's row reader. With Comet on, a generated Spark operator instead reads the - * cached vectors through a ColumnarToRowExec fused into its generated code, which is how these - * reads run when Comet is enabled without native execution. With native execution on, Spark - * operators above the cache read CometInMemoryTableScan's batches through CometColumnarToRow - * instead, which this does not measure. + * shape where Comet's format has the most to lose. Comet's format is read two ways. With Comet + * off, every Spark operator reads rows from the cache scan's row reader. With Comet on, a + * generated Spark operator instead reads the cached vectors through a ColumnarToRowExec fused + * into its generated code, which is how these reads run when Comet is enabled without native + * execution. With native execution on, Spark operators above the cache read + * CometInMemoryTableScan's batches through CometColumnarToRow instead, which + * runAdaptiveBenchmark measures. * * Both formats are cached from the same relation, one copy at a time as in runCodecBenchmark, * and each case checks which serializer cached the relation it reads and which reader it uses. @@ -377,14 +465,7 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { .createOrReplaceTempView(view) val cache = new OneCachedCopy(view) - Seq( - ("row count only (0 of 6 columns)", s"SELECT count(*) FROM $view", 0), - ("narrow projection (1 of 6 columns)", s"SELECT count(k) FROM $view", 1), - ("3 of 6 columns", s"SELECT sum(id), sum(k), sum(v) FROM $view", 3), - ( - "full projection (6 of 6 columns)", - s"SELECT count(id), count(k), count(v), count(s1), count(s2), count(s3) FROM $view", - 6)).foreach { case (label, query, scanned) => + readShapes(view).foreach { case (label, query, scanned) => val benchmark = new Benchmark( s"in-memory cache read by Spark operators, $label", relation.rows.toLong, @@ -446,17 +527,18 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { fused: Boolean = false) /** - * Holds one cached copy of `view` at a time. Two copies could not coexist anyway: the cache - * manager keys on the plan rather than the name, so a second one would find the first. + * Holds one cached copy of `view` at a time, cached under `conf`. Two copies could not coexist + * anyway: the cache manager keys on the plan rather than the name, so a second one would find + * the first. */ - private class OneCachedCopy(view: String) { + private class OneCachedCopy(view: String, conf: Seq[(String, String)] = sparkOperatorConf) { private var cachedBy: String = _ def cacheBy(serializer: String): Unit = if (cachedBy != serializer) { spark.catalog.uncacheTable(view) cachedBy = null withCacheSerializer(serializer) { - withSQLConf(sparkOperatorConf: _*) { + withSQLConf(conf: _*) { spark.catalog.cacheTable(view) spark.table(view).count() } @@ -488,6 +570,116 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { } } + /** + * What the feature changes for a query that runs with Comet, against Spark's own cache format, + * with AQE on and Comet's other settings at their defaults. This is the comparison an + * application gets from turning the feature on or off. Both formats are cached from the same + * relation, one copy at a time as in runSparkOperatorBenchmark. + * + * Two shapes of plan read the cache. With Comet operators above the cache scan, Comet's format + * runs the whole query natively, while Spark's leaves the operators directly above its scan on + * Spark: Comet reads Spark's cache scan only through spark.comet.convert.inMemoryCache.enabled, + * which is off by default. With a Spark operator above the scan, Comet's format is read by the + * native scan and converted to rows for that operator, where Spark's is read by Spark's own + * scan. The Spark operator is the aggregate, with Comet's turned off, standing in for any + * operator Comet does not support. + */ + private def runAdaptiveBenchmark(relation: CachedRelation): Unit = { + val view = s"${relation.table}_adaptive" + val formats = + Seq("Spark's cache format" -> sparkSerializer, "Comet's cache format" -> cometSerializer) + val operatorsAbove = Seq( + "Comet operators" -> Seq.empty[(String, String)], + "a Spark operator" -> Seq(CometConf.COMET_EXEC_AGGREGATE_ENABLED.key -> "false")) + + spark.catalog.clearCache() + withTempTable(view) { + spark + .sql(s"SELECT ${relation.columns.mkString(", ")} FROM ${relation.source}") + .createOrReplaceTempView(view) + + val cache = new OneCachedCopy(view, adaptiveConf) + for { + (operators, operatorConf) <- operatorsAbove + (label, query, scanned) <- readShapes(view) + } { + val benchmark = new Benchmark( + s"in-memory cache with AQE, $operators above the scan, $label", + relation.rows.toLong, + output = output) + formats.foreach { case (name, serializer) => + var verified = false + // Re-caching in this case's format is setup, so it is outside the timer, and it only + // happens on the case's first call, which is a warmup iteration. + benchmark.addTimerCase(name) { timer => + cache.cacheBy(serializer) + withSQLConf(adaptiveConf ++ operatorConf: _*) { + if (!verified) { + verifyAdaptiveRead(query, scanned, serializer, operatorConf.nonEmpty) + verified = true + } + timer.startTiming() + spark.sql(query).noop() + timer.stopTiming() + } + } + } + benchmark.run() + } + + spark.catalog.uncacheTable(view) + } + } + + // The reads runSparkOperatorBenchmark and runAdaptiveBenchmark measure: no columns of the flat + // relation, one, three, and all six. + private def readShapes(view: String): Seq[(String, String, Int)] = Seq( + ("row count only (0 of 6 columns)", s"SELECT count(*) FROM $view", 0), + ("narrow projection (1 of 6 columns)", s"SELECT count(k) FROM $view", 1), + ("3 of 6 columns", s"SELECT sum(id), sum(k), sum(v) FROM $view", 3), + ( + "full projection (6 of 6 columns)", + s"SELECT count(id), count(k), count(v), count(s1), count(s2), count(s3) FROM $view", + 6)) + + // Pins what an adaptive case claims, in the plan AQE settles on, which it does only by running + // the query: one cache scan, native exactly when it reads Comet's format, reading the columns its + // label counts from a relation the named serializer cached. Spark aggregates run only where the + // case puts them: above Spark's scan, which nothing bridges into Comet, or wherever Comet's + // aggregate is turned off. + private def verifyAdaptiveRead( + query: String, + scanned: Int, + serializer: String, + sparkOperator: Boolean): Unit = { + val df = spark.sql(query) + df.collect() + val executed = df.queryExecution.executedPlan + val plan = executed.toString() + assert(executed.isInstanceOf[AdaptiveSparkPlanExec], s"Expected an adaptive plan:\n$plan") + + val nativeScans = collect(executed) { case s: CometInMemoryTableScanExec => s } + val sparkScans = collect(executed) { case s: InMemoryTableScanExec => s } + assert( + nativeScans.length + sparkScans.length == 1, + s"Expected exactly one cache scan:\n$plan") + val cometFormat = serializer == cometSerializer + assert( + nativeScans.nonEmpty == cometFormat, + s"Expected a native scan exactly for Comet's format:\n$plan") + val (relation, columns) = nativeScans.headOption + .map(s => (s.originalPlan.relation, s.scanOutput.length)) + .getOrElse((sparkScans.head.relation, sparkScans.head.attributes.length)) + assert(columns == scanned, s"Expected the scan to read $scanned columns:\n$plan") + val actual = relation.cacheBuilder.serializer.getClass.getName + assert(actual == serializer, s"Expected a relation cached by $serializer, not $actual") + + val sparkAggregates = collect(executed) { case a: HashAggregateExec => a } + assert( + sparkAggregates.nonEmpty == (sparkOperator || !cometFormat), + s"Expected Spark aggregates only where this case puts them:\n$plan") + } + // spark.sql.cache.serializer is static, and InMemoryRelation memoizes the serializer it names // for the life of the JVM. It looks the name up in the active session's conf when a relation is // cached, though, so setting it there directly and clearing the memoized instance around one @@ -535,20 +727,26 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { /** What the cached relation behind `view` occupies, summed over its batches as written. */ private def cachedBytes(view: String): Long = { - val relation = spark - .table(view) - .queryExecution - .optimizedPlan - .collectFirst { case r: InMemoryRelation => r } - .getOrElse(sys.error(s"$view is not cached")) + val relation = cachedRelation(view) // The stored payloads rather than computeStats, which reports the relation's decoded size and // so is the same for every codec. assert(relation.cacheBuilder.isCachedColumnBuffersLoaded, s"$view is not materialized") relation.cacheBuilder.cachedColumnBuffers - .map(CometCachedBatchHelper.payloadSize) + .map { + case batch: DefaultCachedBatch => batch.buffers.map(_.length.toLong).sum + case batch => CometCachedBatchHelper.payloadSize(batch) + } .fold(0L)(_ + _) } + private def cachedRelation(view: String): InMemoryRelation = + spark + .table(view) + .queryExecution + .optimizedPlan + .collectFirst { case r: InMemoryRelation => r } + .getOrElse(sys.error(s"$view is not cached")) + private def runStatsBenchmark(): Unit = { val batchSize = 10000 val types: Seq[DataType] = Seq.fill(3)(LongType) ++ Seq.fill(3)(StringType) @@ -629,8 +827,8 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { // scan and convert" -- which is the overhead this feature exists to remove. // // Neither case is a baseline for Spark's own cache format: both read the same Comet-written - // CometCachedBatch. Spark's format is only measured by runSparkOperatorBenchmark, with Comet - // off, since that is the only comparison it answers. + // CometCachedBatch. Spark's format is measured by runSparkOperatorBenchmark, with Comet off, + // and by runAdaptiveBenchmark, with Comet on. withSQLConf(cacheConf(nativeCacheEnabled = true): _*) { spark .sql(s"SELECT ${relation.columns.mkString(", ")} FROM ${relation.source}") @@ -713,6 +911,16 @@ object CometInMemoryCacheBenchmark extends CometBenchmarkBase { "spark.comet.exec.onHeap.enabled" -> "true", "spark.sql.inMemoryColumnarStorage.batchSize" -> "10000") + // Comet and AQE on, and Comet's other settings at their defaults, unlike cacheConf. The batch + // size matches the other confs, and on-heap mode is what lets Comet run in this session. + private val adaptiveConf: Seq[(String, String)] = Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + "spark.comet.exec.onHeap.enabled" -> "true", + "spark.sql.inMemoryColumnarStorage.batchSize" -> "10000") + private val sparkSerializer = classOf[DefaultCachedBatchSerializer].getName private val cometSerializer = classOf[ArrowCachedBatchSerializer].getName } diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometScalaUDFDispatchBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometScalaUDFDispatchBenchmark.scala new file mode 100644 index 00000000000..3bb3076016f --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometScalaUDFDispatchBenchmark.scala @@ -0,0 +1,207 @@ +/* + * 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.arrow.vector.{BigIntVector, ValueVector, VarBinaryVector} +import org.apache.spark.SparkEnv +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.catalyst.expressions.{Alias, AttributeReference, AttributeSeq, BindReferences, Expression} +import org.apache.spark.sql.execution.adaptive.QueryStageExec + +import org.apache.comet.{CometArrowAllocator, CometConf} +import org.apache.comet.codegen.CometBatchKernelCodegen +import org.apache.comet.codegen.CometBatchKernelCodegen.ArrowColumnSpec +import org.apache.comet.udf.codegen.CometScalaUDFCodegen + +/** + * What a row-level Scala UDF costs on the codegen dispatcher, split into the cost paid per batch + * and the cost paid per row. + * + * Part 1 runs `SELECT max(f(c))` over a Parquet `bigint` column end to end at several batch + * sizes. The total of a cost paid once per batch grows as batches shrink, and the total of a cost + * paid per row does not. A tenth of the rows are null. Spark wraps a UDF with a primitive + * parameter in a null check, `if(isnull(c), null, f(knownnotnull(c)))`, which runs in the + * dispatcher's kernel with the call; a UDF with a boxed parameter gets no such wrapper and + * converts its input through Spark's encoder instead. + * + * Part 2 calls the dispatcher directly on the same in-memory batch over and over, with no Spark + * query and no native code, and does the same with the generated kernel it runs. The difference + * between the two is the dispatcher's own work for each batch: reading its arguments and finding + * the kernel. + * + * To run this benchmark: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometScalaUDFDispatchBenchmark + * }}} + * Results will be written to "spark/benchmarks/CometScalaUDFDispatchBenchmark-**results.txt". + */ +object CometScalaUDFDispatchBenchmark extends CometBenchmarkBase { + + private val Rows = 4L * 1024 * 1024 + private val BatchSizes = Seq(1024, 8192, 65536) + private val LayerBatchRows = Seq(1024, 8192) + + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + spark.udf.register("prim_add_one", (x: Long) => x + 1) + spark.udf.register( + "boxed_add_one", + (x: java.lang.Long) => if (x == null) null else java.lang.Long.valueOf(x + 1)) + + withTempPath { dir => + withTempTable(tbl, "parquetV1Table") { + spark.range(Rows).createOrReplaceTempView(tbl) + prepareTable(dir, spark.sql(s"SELECT IF(id % 10 = 0, NULL, id * 7919) AS c FROM $tbl")) + + checkPlans() + BatchSizes.foreach(endToEnd) + LayerBatchRows.foreach(jvmLayers) + } + } + } + + /** `dispatch` turns the dispatcher on; the cases without a UDF leave it off. */ + private case class EndToEndCase(label: String, f: String, dispatch: Boolean) + + private val endToEndCases = Seq( + EndToEndCase("dispatch, primitive param", "prim_add_one(c)", dispatch = true), + EndToEndCase("dispatch, boxed param", "boxed_add_one(c)", dispatch = true), + EndToEndCase("native c + 1", "c + 1", dispatch = false), + EndToEndCase("no function, max(c)", "c", dispatch = false)) + + private def configs(dispatch: Boolean, batchSize: Int): Seq[(String, String)] = Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_BATCH_SIZE.key -> batchSize.toString, + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> dispatch.toString) + + private def query(f: String): String = s"SELECT max($f) FROM parquetV1Table" + + /** Fails unless Comet runs every operator of every case. */ + private def checkPlans(): Unit = endToEndCases.foreach { + case EndToEndCase(label, f, dispatch) => + withSQLConf(configs(dispatch, BatchSizes.head): _*) { + val df = spark.sql(query(f)) + df.noop() + val plan = stripAQEPlan(df.queryExecution.executedPlan) + // A query stage is a leaf, so check the plan inside each one separately. + val stagePlans = collect(plan) { case s: QueryStageExec => s.plan } + (plan +: stagePlans) + .flatMap(findFirstNonCometOperator(_, classOf[QueryStageExec])) + .headOption + .foreach { op => + throw new IllegalStateException(s"$label ran ${op.nodeName} in Spark:\n$plan") + } + } + } + + private def endToEnd(batchSize: Int): Unit = { + val benchmark = + new Benchmark(s"max(f(c)), batch size $batchSize", Rows, minNumIters = 5, output = output) + endToEndCases.foreach { case EndToEndCase(label, f, dispatch) => + benchmark.addCase(label) { _ => + withSQLConf(configs(dispatch, batchSize): _*) { + spark.sql(query(f)).noop() + } + } + } + benchmark.run() + } + + /** + * The bound subtree the serde would dispatch for `f`, and its closure-serialized bytes. For a + * primitive parameter that is Spark's null check together with the call. + */ + private def boundUdf(f: String): (Expression, Array[Byte]) = { + val plan = spark.sql(s"SELECT $f FROM parquetV1Table").queryExecution.optimizedPlan + val dispatched = plan.expressions.collectFirst { case Alias(child, _) => child }.get + val attrs = dispatched.collect { case a: AttributeReference => a }.distinct + val bound = BindReferences.bindReference(dispatched, AttributeSeq(attrs)) + val buffer = SparkEnv.get.closureSerializer.newInstance().serialize(bound) + val bytes = new Array[Byte](buffer.remaining()) + buffer.get(bytes) + (bound, bytes) + } + + private def binaryScalar(value: Array[Byte]): VarBinaryVector = { + val v = new VarBinaryVector("arg", CometArrowAllocator) + v.allocateNew() + v.setSafe(0, value) + v.setValueCount(1) + v + } + + /** The arguments the serde ships ahead of the data columns. */ + private def dispatcherArgs(bytes: Array[Byte]): Seq[VarBinaryVector] = + Seq(CometScalaUDFCodegen.digest(bytes), bytes).map(binaryScalar) + + private def jvmLayers(batchRows: Int): Unit = { + val in = new BigIntVector("c", CometArrowAllocator) + in.allocateNew(batchRows) + var i = 0 + while (i < batchRows) { + if (i % 10 == 0) in.setNull(i) else in.set(i, i * 7919L) + i += 1 + } + in.setValueCount(batchRows) + val inputs: Array[ValueVector] = Array(in) + val batchesPerIter = (Rows / batchRows).toInt + + val benchmark = + new Benchmark(s"JVM layers, $batchRows-row batches", Rows, minNumIters = 5, output = output) + + def addBatchCase(label: String)(evaluate: () => ValueVector): Unit = + benchmark.addCase(label) { _ => + var b = 0 + while (b < batchesPerIter) { + evaluate().close() + b += 1 + } + } + + val toClose = Seq("prim_add_one(c)" -> "primitive param", "boxed_add_one(c)" -> "boxed param") + .flatMap { case (f, param) => + val (bound, bytes) = boundUdf(f) + val specs = IndexedSeq(ArrowColumnSpec(classOf[BigIntVector], nullable = true)) + val kernel = CometBatchKernelCodegen.compile(bound, specs).newInstance() + kernel.init(0) + val field = CometBatchKernelCodegen.toFfiArrowField("r", bound.dataType, bound.nullable) + addBatchCase(s"generated kernel only, $param") { () => + val out = CometBatchKernelCodegen.allocateOutput(field, batchRows, 0) + kernel.process(inputs, out, batchRows) + out.setValueCount(batchRows) + out + } + + val leading = dispatcherArgs(bytes) + val dispatcher = new CometScalaUDFCodegen + val dispatchInputs: Array[ValueVector] = (leading :+ in).toArray + addBatchCase(s"dispatcher evaluate, $param (${bytes.length}-byte expression)") { () => + dispatcher.evaluate(dispatchInputs, batchRows) + } + leading + } + + try benchmark.run() + finally { + toClose.foreach(_.close()) + in.close() + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleInputConversionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleInputConversionBenchmark.scala new file mode 100644 index 00000000000..3942f1859ef --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleInputConversionBenchmark.scala @@ -0,0 +1,192 @@ +/* + * 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.{CometPlan, CometSparkToColumnarExec} +import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.functions.{col, count, length, lit, sum} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{DataTypes, DecimalType, IntegerType, LongType, StructType} + +import org.apache.comet.CometConf + +// Top-level, so the encoder needs no outer pointer. +case class ShuffleInputBenchRec(a: Long, b: String) + +/** + * Compares three ways to run a shuffle whose input comes from a Spark operator: + * + * - Spark: Comet disabled. + * - Comet: the default. Comet's JVM columnar shuffle buffers the rows and converts them to + * Arrow natively, one partition at a time. + * - Comet, converted: `spark.comet.convert.shuffleInput.enabled`, which converts the rows with + * `CometSparkToColumnarExec` and uses native shuffle. + * + * Each case partitions its rows into 200 partitions and aggregates them. Most hash-partition on a + * number, because the conversion leaves a shuffle that hashes a string on the JVM columnar + * shuffle, and one range-partitions. Every arm's result and plan are checked before it is timed, + * and the Comet 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.CometShuffleInputConversionBenchmark + * }}} + * Results will be written to + * "spark/benchmarks/CometShuffleInputConversionBenchmark-**results.txt". + */ +object CometShuffleInputConversionBenchmark extends CometBenchmarkBase { + + private val numRows = 4L * 1024 * 1024 + private val numPartitions = 200 + + import spark.implicits._ + + private case class Arm(name: String, confs: Seq[(String, String)], check: SparkPlan => Unit) + + // No AQE, so every arm plans the same shape on every iteration. + private val planConfs = Seq(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") + + private def cometConfs(convert: Boolean) = planConfs ++ Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + // Spark scans the RDDs, as it does by default, so the shuffle reads a Spark operator. + CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.key -> convert.toString) + + private def shuffles(plan: SparkPlan): Seq[CometShuffleExchangeExec] = + collect(plan) { case s: CometShuffleExchangeExec => s } + + private val arms = Seq( + Arm( + "Spark", + planConfs :+ (CometConf.COMET_ENABLED.key -> "false"), + plan => assert(collect(plan) { case c: CometPlan => c }.isEmpty, plan)), + Arm( + "Comet", + cometConfs(convert = false), + plan => assert(shuffles(plan).exists(_.shuffleType == CometColumnarShuffle), plan)), + Arm( + "Comet, converted", + cometConfs(convert = true), + plan => + assert( + shuffles(plan).exists(s => + s.shuffleType == CometNativeShuffle && s.child + .isInstanceOf[CometSparkToColumnarExec]), + plan))) + + private val rowSchema = new StructType() + .add("k", IntegerType) + .add("l", LongType) + .add("s", DataTypes.StringType) + .add("d", DataTypes.DoubleType) + .add("m", DecimalType(18, 2)) + + /** `numRows` rows from an RDD, generated as the query runs. */ + private def rddRows(): DataFrame = { + val rows = spark.sparkContext.range(0, numRows, 1, 1).map { id => + val i = id.toInt + Row( + i % 1000, + id, + s"value-${i * 31 % 100003}", + i + 0.5d, + java.math.BigDecimal.valueOf(i * 7919L % 1000000000L, 2)) + } + spark.createDataFrame(rows, rowSchema) + } + + /** Totals the per-group results in `s`, so each case returns one row. */ + private def total(grouped: DataFrame): DataFrame = grouped.agg(sum("s"), count(lit(1))) + + private val cases: Seq[(String, () => DataFrame)] = Seq( + "RDD rows: int, long, double" -> (() => + total( + rddRows() + .select("k", "l", "d") + .repartition(numPartitions, col("k")) + .groupBy("k") + .agg((sum("l") + sum("d")).as("s")))), + "RDD rows: int, long, string, double, decimal(18,2)" -> (() => + total( + rddRows() + .repartition(numPartitions, col("k")) + .groupBy("k") + .agg((sum("l") + sum(length(col("s"))) + sum("d") + sum("m")).as("s")))), + "RDD rows: int, long, string, double, decimal(18,2), range partitioned" -> (() => + total( + rddRows() + .repartitionByRange(numPartitions, col("l")) + .select((col("k") + col("l") + length(col("s")) + col("d").cast("long") + col("m")) + .as("s")))), + "map over a Parquet scan" -> (() => + total( + spark + .table("parquetV1Table") + .select(col("id").as("a"), col("key").as("b")) + .as[ShuffleInputBenchRec] + .map(r => ShuffleInputBenchRec(r.a % 1000, r.b)) + .repartition(numPartitions, col("a")) + .groupBy("a") + .agg(sum(length(col("b"))).as("s"))))) + + 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 = { + withTempPath { dir => + withTempTable("parquetV1Table") { + prepareTable( + dir, + spark.range(numRows).selectExpr("id", "CAST(id % 1000 AS STRING) AS key")) + + 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(1).copy(name = "Comet (repeat)")).foreach { arm => + benchmark.addCase(arm.name) { _ => + withSQLConf(arm.confs: _*) { + query().collect() + } + } + } + benchmark.run() + } + } + } + } + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeWriteSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeWriteSuite.scala new file mode 100644 index 00000000000..098b1c16bc5 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeWriteSuite.scala @@ -0,0 +1,390 @@ +/* + * 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.io.{File, IOException} +import java.util.concurrent.ConcurrentLinkedQueue + +import scala.jdk.CollectionConverters._ +import scala.util.Using + +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.Path +import org.apache.hadoop.mapreduce.{JobContext, TaskAttemptContext} +import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat +import org.apache.parquet.hadoop.{ParquetFileReader, ParquetOutputFormat} +import org.apache.parquet.hadoop.metadata.CompressionCodecName +import org.apache.parquet.hadoop.util.HadoopInputFile +import org.apache.spark.TaskContext +import org.apache.spark.internal.io.FileCommitProtocol.TaskCommitMessage +import org.apache.spark.internal.io.FileNameSpec +import org.apache.spark.sql.{CometTestBase, DataFrame, SaveMode} +import org.apache.spark.sql.execution.command.DataWritingCommandExec +import org.apache.spark.sql.execution.datasources.{InsertIntoHadoopFsRelationCommand, SQLHadoopMapReduceCommitProtocol} +import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat +import org.apache.spark.sql.execution.metric.SQLMetric +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{StringType, StructField} +import org.apache.spark.util.SerializableConfiguration + +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus +import org.apache.comet.serde.{Incompatible, Unsupported} +import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.operator.{schema2Proto, CometDataWritingCommand, NativeWriteUtils} + +/** Exercises the retained Spark 3.x writer, including both terminal execution entry points. */ +class CometNativeWriteSuite extends CometTestBase { + import NativeWriteCommitProtocol._ + + private def withWriter(partitions: Int = 1)( + f: (CometNativeWriteExec, DataFrame, File) => Unit): Unit = { + assume(!isSpark40Plus, "Spark 4.0+ uses CometWriteFilesExec") + withTempPath { dir => + val source = new File(dir, "source").getAbsolutePath + val output = new File(dir, "output with % space") + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + spark.range(0, 20, 1, partitions).write.parquet(source) + } + withSQLConf( + CometConf.COMET_NATIVE_PARQUET_WRITE_ENABLED.key -> "true", + CometConf.COMET_OPERATOR_DATA_WRITING_COMMAND_ALLOW_INCOMPAT.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.FILE_COMMIT_PROTOCOL_CLASS.key -> classOf[NativeWriteCommitProtocol].getName, + CometConf.COMET_METRICS_UPDATE_INTERVAL.key -> "0") { + val data = spark.read.parquet(source) + val sourcePlan = data.queryExecution.executedPlan + val child = sourcePlan + .collectFirst { case scan: CometNativeScanExec => + scan + } + .getOrElse(fail(s"Expected a native source scan: $sourcePlan")) + val cmd = InsertIntoHadoopFsRelationCommand( + outputPath = new Path(output.toURI), + staticPartitions = Map.empty, + ifPartitionNotExists = false, + partitionColumns = Seq.empty, + bucketSpec = None, + fileFormat = new ParquetFileFormat, + options = Map("compression" -> "gzip", "comet.test.option" -> "present"), + query = data.queryExecution.analyzed, + mode = SaveMode.ErrorIfExists, + catalogTable = None, + fileIndex = None, + outputColumnNames = data.columns.toSeq) + val command = DataWritingCommandExec(cmd, child) + val nativeOp = CometDataWritingCommand.convert(command, Operator.newBuilder()).get + val writer = CometDataWritingCommand + .createExec(nativeOp, command) + .asInstanceOf[CometNativeWriteExec] + assert(writer.committer.isInstanceOf[NativeWriteCommitProtocol]) + reset() + try f(writer, data, output) + finally reset() + } + } + } + + test("Spark 3.x HDFS admission ignores the configurable output basename") { + assume(!isSpark40Plus, "Spark 4.0+ honors mapreduce.output.basename") + val data = spark.range(1) + val cmd = InsertIntoHadoopFsRelationCommand( + outputPath = new Path("hdfs://localhost:8020/output"), + staticPartitions = Map.empty, + ifPartitionNotExists = false, + partitionColumns = Seq.empty, + bucketSpec = None, + fileFormat = new ParquetFileFormat, + options = Map(NativeWriteUtils.BASE_OUTPUT_NAME -> "out%?#"), + query = data.queryExecution.analyzed, + mode = SaveMode.ErrorIfExists, + catalogTable = None, + fileIndex = None, + outputColumnNames = data.columns.toSeq) + val command = DataWritingCommandExec(cmd, data.queryExecution.executedPlan) + assert(CometDataWritingCommand.getSupportLevel(command).isInstanceOf[Incompatible]) + val escapedPath = cmd.copy(outputPath = new Path("hdfs://localhost:8020/out%?#")) + assert( + CometDataWritingCommand + .getSupportLevel(command.copy(cmd = escapedPath)) + .isInstanceOf[Unsupported]) + } + + Seq(false, true).foreach { columnar => + test( + s"configured commit protocol completes the ${if (columnar) "columnar" else "row"} write") { + withWriter(2) { (writer, data, output) => + val resultCount = + if (columnar) writer.executeColumnar().count() else writer.execute().count() + assert(resultCount == 0) + assert(events.asScala.head == "setupJob") + assert(events.asScala.last == "commitJob") + assert(!events.asScala.exists(_.startsWith("abort"))) + assert(writer.metrics("rows_written").value == 20, "the writer must execute natively") + val files = allocated.asScala.map(p => new Path(p).getName).toSet + assert(files.nonEmpty) + assert(files.forall(n => n.startsWith("chosen % ") && n.endsWith(".gz.parquet"))) + assert(output.list().filter(_.endsWith(".parquet")).toSet == files) + assert(new File(output, "_SUCCESS").isFile) + assert(!new File(output, "_temporary").exists()) + checkAnswer(spark.read.parquet(output.getAbsolutePath), data) + } + } + } + + test("the footer codec matches the codec in the file name") { + withWriter() { (writer, data, output) => + // The planner chose gzip, but the job configuration says zstd. The job configuration also + // names the file. + val conf = new Configuration(writer.serializableHadoopConf.value) + conf.set(ParquetOutputFormat.COMPRESSION, CompressionCodecName.ZSTD.name()) + writer.copy(serializableHadoopConf = new SerializableConfiguration(conf)).execute().count() + + val files = output.listFiles().filter(_.getName.endsWith(".parquet")) + assert(files.nonEmpty) + files.foreach { file => + assert(file.getName.endsWith(".zstd.parquet"), s"Expected a zstd file: ${file.getName}") + val input = HadoopInputFile.fromPath(new Path(file.toURI), conf) + Using.resource(ParquetFileReader.open(input)) { reader => + val codecs = reader.getFooter.getBlocks.asScala + .flatMap(_.getColumns.asScala) + .map(_.getCodec) + .toSet + assert(codecs == Set(CompressionCodecName.ZSTD), s"${file.getName} uses $codecs") + } + } + checkAnswer(spark.read.parquet(output.getAbsolutePath), data) + } + } + + test("native iterator creation failure aborts the task and job") { + withWriter() { (writer, _, output) => + // Inject an unknown native config after the child scan has initialized, so the writer's + // CometExecIterator constructor fails in createPlan, before the drain/close block is entered. + failAt = "createIterator" + failAbort = true + val error = intercept[Exception](writer.execute().count()) + val causes = allErrors(error) + assert(causes.exists(_.getStackTrace.exists(_.getMethodName == "createPlan"))) + assert(causes.exists(e => Option(e.getMessage).exists(_.contains("no_such_namespace")))) + assert(causes.exists(_.getSuppressed.exists(_.getMessage == "injected abortTask failure"))) + assert(error.getSuppressed.exists(_.getMessage == "injected abortJob failure")) + assertAborted(output) + assert(!events.contains("commitTask")) + } + } + + test("native execution failure survives throwing cleanup and abort callbacks") { + withWriter() { (writer, _, output) => + // A mismatched output type fails when writing the first native batch, after the native + // stream is initialized. This also lets releasePlan exercise its final metrics callback. + val badWriter = writer.nativeOp.getParquetWriter.toBuilder + .clearOutputSchema() + .addAllOutputSchema(schema2Proto(Seq(StructField("id", StringType))).asJava) + .build() + val badPlan = writer.nativeOp.toBuilder.setParquetWriter(badWriter).build() + failAbort = true + val failingWriter = new CleanupFailingNativeWriteExec(writer.copy(nativeOp = badPlan)) + val error = intercept[Exception](failingWriter.executeColumnar().count()) + val causes = allErrors(error) + assert(causes.exists(_.getStackTrace.exists(_.getMethodName == "executePlan"))) + assert(causes.exists(e => + Option(e.getMessage).exists(_.contains("Failed to rename batch columns")))) + assert(causes.exists(_.getSuppressed.exists(_.getMessage == "injected cleanup failure"))) + assert(causes.exists(_.getSuppressed.exists(_.getMessage == "injected abortTask failure"))) + assert(error.getSuppressed.exists(_.getMessage == "injected abortJob failure")) + assertAborted(output) + assert(!events.contains("commitTask")) + } + } + + test("a cleanup failure aborts instead of committing") { + withWriter() { (writer, _, output) => + val error = + intercept[Exception](new CleanupFailingNativeWriteExec(writer).execute().count()) + assert(allErrors(error).exists(_.getMessage == "injected cleanup failure")) + assertAborted(output) + assert(!events.contains("commitTask")) + } + } + + Seq("commitTask", "onTaskCommit", "commitJob").foreach { phase => + test(s"$phase failure aborts and preserves the commit failure") { + withWriter() { (writer, _, output) => + failAt = phase + failAbort = true + val error = intercept[Exception](writer.execute().count()) + assert(allErrors(error).exists(_.getMessage == s"injected $phase failure")) + assert(error.getSuppressed.exists(_.getMessage == "injected abortJob failure")) + assert(events.contains("abortJob")) + if (phase == "commitTask") { + assert(events.contains("abortTask")) + assert( + allErrors(error).exists( + _.getSuppressed.exists(_.getMessage == "injected abortTask failure"))) + } + assert(!new File(output, "_SUCCESS").exists()) + assert(!new File(output, "_temporary").exists()) + assert( + written.asScala.nonEmpty, + "the failure must occur after the native file was written") + } + } + } + + private def allErrors(error: Throwable): Seq[Throwable] = + Seq(error) ++ Option(error.getCause).toSeq.flatMap(allErrors) ++ + error.getSuppressed.toSeq.flatMap(allErrors) + + private def assertAborted(output: File): Unit = { + assert(events.contains("abortTask")) + assert(events.contains("abortJob")) + assert(!events.contains("commitJob")) + assert(!new File(output, "_SUCCESS").exists()) + assert(!new File(output, "_temporary").exists()) + assert(allocated.asScala.nonEmpty, "the test must reach task file allocation") + } +} + +/** + * Local-mode probes are reset for each write; executor and driver callbacks share this object. + */ +object NativeWriteCommitProtocol { + val events = new ConcurrentLinkedQueue[String]() + val allocated = new ConcurrentLinkedQueue[String]() + val written = new ConcurrentLinkedQueue[String]() + val received = new ConcurrentLinkedQueue[TaskCommitMessage]() + @volatile var failAt: String = "" + @volatile var failAbort = false + + def reset(): Unit = { + events.clear() + allocated.clear() + written.clear() + received.clear() + failAt = "" + failAbort = false + } +} + +class NativeWriteCommitProtocol(jobId: String, path: String, dynamicPartitionOverwrite: Boolean) + extends SQLHadoopMapReduceCommitProtocol(jobId, path, dynamicPartitionOverwrite) { + import NativeWriteCommitProtocol._ + private var taskFile: String = _ + + override def setupJob(context: JobContext): Unit = { + assert(context.getConfiguration.get("comet.test.option") == "present") + assert(context.getOutputFormatClass == classOf[ParquetOutputFormat[_]]) + assert(FileOutputFormat.getOutputPath(context) == new Path(path)) + assert(context.getConfiguration.get("spark.sql.sources.writeJobUUID") != null) + super.setupJob(context) + context.getConfiguration.set("comet.test.setupJob", "present") + events.add("setupJob") + } + + override def setupTask(context: TaskAttemptContext): Unit = { + assert(context.getConfiguration.get("comet.test.option") == "present") + assert(context.getConfiguration.get("comet.test.setupJob") == "present") + super.setupTask(context) + events.add("setupTask") + } + + override def newTaskTempFile( + context: TaskAttemptContext, + dir: Option[String], + spec: FileNameSpec): String = { + val sparkPath = new Path(super.newTaskTempFile(context, dir, spec)) + taskFile = new Path(sparkPath.getParent, "chosen % " + sparkPath.getName).toString + allocated.add(taskFile) + if (failAt == "createIterator") { + val properties = TaskContext.get().getLocalProperties + properties.setProperty(CometConf.COMET_RESPECT_DATAFUSION_CONFIGS.key, "true") + properties.setProperty("spark.comet.datafusion.no_such_namespace.option", "1") + } + taskFile + } + + override def commitTask(context: TaskAttemptContext): TaskCommitMessage = { + val file = new Path(taskFile) + assert(file.getFileSystem(context.getConfiguration).getFileStatus(file).getLen > 0) + written.add(taskFile) + events.add("commitTask") + if (failAt == "commitTask") throw new IOException("injected commitTask failure") + new TaskCommitMessage((context.getTaskAttemptID.toString, super.commitTask(context))) + } + + override def onTaskCommit(message: TaskCommitMessage): Unit = { + events.add("onTaskCommit") + received.add(message) + if (failAt == "onTaskCommit") throw new IOException("injected onTaskCommit failure") + super.onTaskCommit(message.obj.asInstanceOf[(String, TaskCommitMessage)]._2) + } + + override def commitJob(context: JobContext, messages: Seq[TaskCommitMessage]): Unit = { + assert(messages.nonEmpty) + assert(messages.toSet == received.asScala.toSet) + assert( + messages.map(_.obj.asInstanceOf[(String, TaskCommitMessage)]._1).distinct.size == + messages.size) + events.add("commitJob") + if (failAt == "commitJob") throw new IOException("injected commitJob failure") + super.commitJob(context, messages.map(_.obj.asInstanceOf[(String, TaskCommitMessage)]._2)) + } + + override def abortTask(context: TaskAttemptContext): Unit = { + events.add("abortTask") + super.abortTask(context) + if (failAbort) throw new IOException("injected abortTask failure") + } + + override def abortJob(context: JobContext): Unit = { + events.add("abortJob") + super.abortJob(context) + if (failAbort) throw new IOException("injected abortJob failure") + } +} + +/** + * With periodic updates disabled, this metric throws only during native iterator teardown. + * CometMetricNode reports a sum metric with `add`, so both update methods throw. + */ +class CleanupFailingWriteMetric extends SQLMetric("sum") { + override def copy(): SQLMetric = new CleanupFailingWriteMetric + override def set(value: Long): Unit = throw new IOException("injected cleanup failure") + override def add(value: Long): Unit = throw new IOException("injected cleanup failure") +} + +class CleanupFailingNativeWriteExec(writer: CometNativeWriteExec) + extends CometNativeWriteExec( + writer.nativeOp, + writer.originalPlan, + writer.child, + writer.outputPath, + writer.mode, + writer.committer, + writer.serializableHadoopConf, + writer.outputWriterFactory, + writer.jobTrackerID) { + override lazy val metrics: Map[String, SQLMetric] = { + val failingMetric = new CleanupFailingWriteMetric + sparkContext.register(failingMetric) + writer.metrics.updated("rows_written", failingMetric) + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/DecimalPrecisionSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/DecimalPrecisionSuite.scala new file mode 100644 index 00000000000..be4c7db5de8 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/DecimalPrecisionSuite.scala @@ -0,0 +1,83 @@ +/* + * 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.SparkFunSuite +import org.apache.spark.sql.catalyst.expressions.{Add, Alias, AttributeReference, CaseWhen, Cast, CheckOverflow, Coalesce, Divide, Expression, GreaterThan, IsNotNull, Literal, Multiply, Remainder, Subtract} +import org.apache.spark.sql.types.{DecimalType, IntegerType, LongType, StringType} + +class DecimalPrecisionSuite extends SparkFunSuite { + + private val decA = AttributeReference("a", DecimalType(10, 2))() + private val decB = AttributeReference("b", DecimalType(10, 2))() + private val intA = AttributeReference("c", IntegerType)() + private val intB = AttributeReference("d", IntegerType)() + + private def assertUnchanged(expr: Expression): Unit = { + // Reference equality, not `===`: the point of the guard is that no node is rebuilt. + val promoted = DecimalPrecision.promote(expr) + assert(promoted.eq(expr), s"expected promote to leave $expr untouched, got: $promoted") + } + + test("promote leaves trees without decimal arithmetic untouched") { + Seq( + decA, + Literal(1), + Add(intA, intB), + Multiply(intA, Cast(decA, IntegerType)), + GreaterThan(decA, decB), + Cast(Add(intA, intB), DecimalType(20, 4)), + Coalesce(Seq(decA, decB)), + CaseWhen(Seq((IsNotNull(decA), decA)), Some(decB)), + Alias(Add(Cast(decA, LongType), Cast(decB, LongType)), "x")()).foreach(assertUnchanged) + } + + test("promote wraps decimal arithmetic wherever it appears in the tree") { + Seq[Expression]( + Add(decA, decB), + Subtract(decA, decB), + Multiply(decA, decB), + Divide(decA, decB), + Remainder(decA, decB)).foreach { arithmetic => + DecimalPrecision.promote(arithmetic) match { + case CheckOverflow(child, dt, _) => + assert(child.eq(arithmetic)) + assert(dt === arithmetic.dataType) + case other => fail(s"expected $arithmetic to be wrapped in CheckOverflow, got: $other") + } + + // The guard has to inspect the whole tree, not just the root: a decimal arithmetic node + // buried under expressions the rule does not rewrite must still be promoted. + val buried = Cast(Coalesce(Seq(decA, arithmetic)), StringType) + val promoted = DecimalPrecision.promote(buried) + assert(!promoted.eq(buried)) + assert(promoted.collectFirst { case c: CheckOverflow => c }.isDefined) + } + } + + test("promote rewrites only the decimal arithmetic in a mixed tree") { + val intAdd = Add(intA, intB) + val decAdd = Add(decA, decB) + val promoted = DecimalPrecision.promote(GreaterThan(Cast(intAdd, DecimalType(20, 2)), decAdd)) + val overflows = promoted.collect { case c: CheckOverflow => c } + assert(overflows.length === 1) + assert(overflows.head.child.eq(decAdd)) + } +} diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala index ef73522577a..34bb32d908f 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala @@ -249,7 +249,7 @@ class CometArrowStreamSuite extends AnyFunSuite with Matchers { val allocator = new RootAllocator(Long.MaxValue) val numRows = 32 class CountingFixedWidthWriter(override val valueVector: BaseFixedWidthVector) - extends FixedWidthArrowFieldWriter { + extends FixedWidthArrowFieldWriter(valueVector) { var scalarWrites: Int = 0 override def setValue(input: SpecializedGetters, ordinal: Int): Unit = scalarWrites += 1 override protected def setValueUnsafe(input: SpecializedGetters, ordinal: Int): Unit = diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowWriterSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowWriterSuite.scala index 294f9152985..400d9aad797 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowWriterSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowWriterSuite.scala @@ -19,23 +19,28 @@ package org.apache.spark.sql.comet.execution.arrow +import java.lang.reflect.Modifier import java.math.{BigDecimal => JavaBigDecimal, BigInteger} +import java.util.{ArrayDeque, Collections, IdentityHashMap} +import scala.collection.mutable.ArrayBuffer import scala.jdk.CollectionConverters._ import scala.util.Random import org.scalatest.funsuite.AnyFunSuite import org.scalatest.matchers.should.Matchers -import org.apache.arrow.memory.RootAllocator -import org.apache.arrow.vector.{DecimalVector, FieldVector, IntVector, ValueVector, VarCharVector, VectorSchemaRoot} +import org.apache.arrow.memory.{ArrowBuf, RootAllocator} +import org.apache.arrow.vector.{BaseVariableWidthVector, DecimalVector, FieldVector, IntVector, ValueVector, VarCharVector, VectorSchemaRoot} import org.apache.arrow.vector.complex.{ListVector, StructVector} -import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeProjection} -import org.apache.spark.sql.catalyst.util.GenericArrayData +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeArrayData, UnsafeMapData, UnsafeProjection, UnsafeRow} +import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, GenericArrayData} import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.vectorized.{ConstantColumnVector, Dictionary, OffHeapColumnVector, OnHeapColumnVector, WritableColumnVector} import org.apache.spark.sql.types._ import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} +import org.apache.spark.unsafe.types.UTF8String import org.apache.comet.CometSparkSessionExtensions.isSpark41Plus @@ -124,6 +129,7 @@ class CometArrowWriterSuite extends AnyFunSuite with Matchers { /** * Fills rows `[0, n)`. Array and map rows are laid out back to back, or from the last row * backwards when `reversed`, and null rows get no offsets at all, as in Spark's readers. + * Collections hold fewer than `maxLength` elements. */ private def fill( v: WritableColumnVector, @@ -131,18 +137,19 @@ class CometArrowWriterSuite extends AnyFunSuite with Matchers { n: Int, rnd: Random, nullFraction: Double, - reversed: Boolean): Unit = dataType match { + reversed: Boolean, + maxLength: Int = 6): Unit = dataType match { case st: StructType => (0 until n).foreach(i => if (rnd.nextDouble() < nullFraction) v.putNull(i)) st.fields.zipWithIndex.foreach { case (field, ordinal) => val child = v.getChild(ordinal) child.reserve(n) - fill(child, field.dataType, n, rnd, nullFraction, reversed) + fill(child, field.dataType, n, rnd, nullFraction, reversed, maxLength) // Spark's Parquet reader nulls the fields of a null struct. Other producers need not. (0 until n).foreach(i => if (v.isNullAt(i) && rnd.nextBoolean()) child.putNull(i)) } case _: ArrayType | _: MapType => - val lengths = Array.fill(n)(rnd.nextInt(6)) + val lengths = Array.fill(n)(rnd.nextInt(maxLength)) val children = dataType match { case _: ArrayType => Seq(v.arrayData()) case _ => Seq(v.getChild(0), v.getChild(1)) @@ -159,10 +166,10 @@ class CometArrowWriterSuite extends AnyFunSuite with Matchers { } dataType match { case ArrayType(elementType, _) => - fill(v.arrayData(), elementType, offset, rnd, nullFraction, reversed) + fill(v.arrayData(), elementType, offset, rnd, nullFraction, reversed, maxLength) case MapType(keyType, valueType, _) => - fill(v.getChild(0), keyType, offset, rnd, 0.0, reversed) - fill(v.getChild(1), valueType, offset, rnd, nullFraction, reversed) + fill(v.getChild(0), keyType, offset, rnd, 0.0, reversed, maxLength) + fill(v.getChild(1), valueType, offset, rnd, nullFraction, reversed, maxLength) } case _ => (0 until n).foreach { i => @@ -206,6 +213,14 @@ class CometArrowWriterSuite extends AnyFunSuite with Matchers { } } + /** Both roots must hold the same rows, compared vector by vector. */ + private def assertSameRoots(expected: VectorSchemaRoot, actual: VectorSchemaRoot): Unit = { + actual.getRowCount shouldBe expected.getRowCount + expected.getFieldVectors.asScala.zip(actual.getFieldVectors.asScala).foreach { case (e, a) => + assertSameVectors(e, a, "") + } + } + /** Each vector in the tree must match, leaf values under null parents included. */ private def assertSameVectors( expected: ValueVector, @@ -222,12 +237,20 @@ class CometArrowWriterSuite extends AnyFunSuite with Matchers { a.getElementStartIndex(i) shouldBe e.getElementStartIndex(i) a.getElementEndIndex(i) shouldBe e.getElementEndIndex(i) case (_: StructVector, _) => - case _ if !expected.isNull(i) => - (expected.getObject(i), actual.getObject(i)) match { - case (e: Array[Byte], a: Array[Byte]) => a.toSeq shouldBe e.toSeq - case (e, a) => a shouldBe e - } case _ => + // A reader takes each value's start from the end of the one before, null or not. + // The paths may differ in the bytes under a null struct, but never go backwards. + actual match { + case a: BaseVariableWidthVector => + a.getStartOffset(i + 1) should be >= a.getStartOffset(i) + case _ => + } + if (!expected.isNull(i)) { + (expected.getObject(i), actual.getObject(i)) match { + case (e: Array[Byte], a: Array[Byte]) => a.toSeq shouldBe e.toSeq + case (e, a) => a shouldBe e + } + } } } } @@ -269,9 +292,7 @@ class CometArrowWriterSuite extends AnyFunSuite with Matchers { rowWriter.finish() columnar.getRowCount shouldBe length - rows.getFieldVectors.asScala.zip(columnar.getFieldVectors.asScala).foreach { case (e, a) => - assertSameVectors(e, a, "") - } + assertSameRoots(rows, columnar) } finally { columnar.close() rows.close() @@ -488,6 +509,35 @@ class CometArrowWriterSuite extends AnyFunSuite with Matchers { } } + test("a decimal past its precision in an unsafe array still fails the row path") { + // Written at a wider precision and read at precision p, so the array holds 10^p or -10^p, the + // smallest values past p. Up to 18 digits the array holds the unscaled long, and past that the + // unscaled bytes. + Seq( + (DecimalType(18, 2), DecimalType(5, 2), "1000.00"), + (DecimalType(38, 0), DecimalType(20, 0), "100000000000000000000")).foreach { + case (written, read, magnitude) => + Seq(magnitude, s"-$magnitude").foreach { value => + val decimal = Decimal(new JavaBigDecimal(value), written.precision, written.scale) + val row = UnsafeProjection.create(new StructType().add("a", ArrayType(written)))( + new GenericInternalRow(Array[Any](new GenericArrayData(Array[Any](decimal))))) + val allocator = new RootAllocator(Long.MaxValue) + val root = VectorSchemaRoot.create( + Utils.toArrowSchema(new StructType().add("a", ArrayType(read)), "UTC"), + allocator) + try { + val writer = ArrowWriter.create(root, 1) + withClue(s"$read $value: ") { + intercept[ArithmeticException](writer.write(row)) + } + } finally { + root.close() + allocator.close() + } + } + } + } + test("a narrow decimal with more digits than its precision passes through") { // Spark's getDecimal does not check an int- or long-backed value against the precision, so // neither path does. Arrow's BigDecimal setter, which the writer used before, threw instead. @@ -588,6 +638,52 @@ class CometArrowWriterSuite extends AnyFunSuite with Matchers { } } + test("a dictionary-encoded field appends as a column after its struct took the row path") { + // A struct with a collection field takes the row path in a batch that holds null structs, + // which leaves the offsets of its string field's trailing nulls for the next value to fill. A + // dictionary-encoded field has to fill them, and grow its buffers, before it appends a column. + val st = new StructType().add("a", ArrayType(IntegerType)).add("s", StringType) + val schema = new StructType().add("st", st) + val rnd = new Random(5) + val batches = Seq((0.5, 100), (0.0, 5000), (0.5, 100), (0.0, 5000)).map { + case (nullFraction, n) => + val v = newVector(n, st, offHeap = false) + (0 until n).foreach { i => + if (rnd.nextDouble() < nullFraction || (nullFraction > 0 && i == n - 1)) { + v.putNull(i) + } + } + fill(v.getChild(0), ArrayType(IntegerType), n, rnd, nullFraction = 0.0, reversed = false) + fillDictionary(v.getChild(1), StringType, n, rnd, nullFraction = 0.0) + // Spark's Parquet reader nulls the fields of a null struct. + (0 until n).foreach { i => + if (v.isNullAt(i)) { + v.getChild(0).putNull(i) + v.getChild(1).putNull(i) + } + } + new ColumnarBatch(Array[ColumnVector](v), n) + } + val allocator = new RootAllocator(Long.MaxValue) + val arrowSchema = Utils.toArrowSchema(schema, "UTC") + val columnar = VectorSchemaRoot.create(arrowSchema, allocator) + val rows = VectorSchemaRoot.create(arrowSchema, allocator) + try { + val columnarWriter = ArrowWriter.create(columnar, 1) + batches.foreach(b => columnarWriter.writeColumns(b, 0, b.numRows())) + columnarWriter.finish() + val rowWriter = ArrowWriter.create(rows, 1) + batches.foreach(b => (0 until b.numRows()).foreach(i => rowWriter.write(b.getRow(i)))) + rowWriter.finish() + assertSameRoots(rows, columnar) + } finally { + columnar.close() + rows.close() + allocator.close() + batches.foreach(_.close()) + } + } + test("a struct that switches between the columnar and row paths across appended batches") { // A struct with an array or a map field takes the row path only in batches that hold null // structs, so its fields' writers keep appending where the other path left off. @@ -627,10 +723,7 @@ class CometArrowWriterSuite extends AnyFunSuite with Matchers { val rowWriter = ArrowWriter.create(rows, batches.map(_.numRows()).sum) batches.foreach(b => (0 until b.numRows()).foreach(i => rowWriter.write(b.getRow(i)))) rowWriter.finish() - columnar.getRowCount shouldBe rows.getRowCount - rows.getFieldVectors.asScala.zip(columnar.getFieldVectors.asScala).foreach { - case (e, a) => assertSameVectors(e, a, "") - } + assertSameRoots(rows, columnar) } finally { columnar.close() rows.close() @@ -640,4 +733,221 @@ class CometArrowWriterSuite extends AnyFunSuite with Matchers { } } } + + // Nested shapes whose unsafe forms take paths of their own: an array of each primitive type not + // in `nestedTypes`, whose elements are either copied in one block or converted one at a time, + // collections inside collections, and a struct wider than one word of null bits. + private val moreNestedTypes: Seq[DataType] = + primitiveTypes.map(ArrayType(_)).filterNot(nestedTypes.contains) ++ Seq( + ArrayType(MapType(StringType, IntegerType)), + MapType(StringType, ArrayType(StringType)), + MapType(LongType, DecimalType(9, 2)), + new StructType() + .add("m", MapType(IntegerType, StringType)) + .add("s", new StructType().add("x", BinaryType).add("y", DecimalType(18, 4))), + StructType( + (0 until 70).map(i => StructField(s"f$i", if (i % 7 == 3) StringType else LongType)))) + + /** + * Writes `numRows` rows, `row(i)` through the generic row path and its unsafe projection + * through the unsafe one, the latter copied to off-heap memory when `offHeap`, as Spark's + * off-heap pages hold rows. A third writer alternates between the two kinds of row, so each + * path appends where the other left off, and a fourth writes generic rows holding the unsafe + * row's values, as a copied row holds unsafe arrays. All of them must write the same vectors. + */ + private def assertUnsafeRowsMatchGeneric(numRows: Int, schema: StructType, offHeap: Boolean)( + row: Int => InternalRow): Unit = { + val allocator = new RootAllocator(Long.MaxValue) + val arrowSchema = Utils.toArrowSchema(schema, "UTC") + val roots = Seq.fill(4)(VectorSchemaRoot.create(arrowSchema, allocator)) + val buffers = ArrayBuffer.empty[ArrowBuf] + try { + // Every writer starts undersized, so each buffer has to grow. + val writers = roots.map(ArrowWriter.create(_, 1)) + val project = UnsafeProjection.create(schema) + (0 until numRows).foreach { i => + val generic = row(i) + val projected = project(generic) + val unsafe = if (offHeap) { + val buffer = allocator.buffer(math.max(projected.getSizeInBytes, 8).toLong) + buffers += buffer + projected.writeToMemory(null, buffer.memoryAddress()) + val copy = new UnsafeRow(projected.numFields()) + copy.pointTo(null, buffer.memoryAddress(), projected.getSizeInBytes) + copy + } else { + projected + } + writers(0).write(generic) + writers(1).write(unsafe) + writers(2).write(if (i % 3 == 1) generic else unsafe) + writers(3).write(new GenericInternalRow(Array.tabulate[Any](schema.length) { ordinal => + unsafe.get(ordinal, schema(ordinal).dataType) + })) + } + writers.foreach(_.finish()) + roots.head.getRowCount shouldBe numRows + roots.tail.foreach(assertSameRoots(roots.head, _)) + } finally { + roots.foreach(_.close()) + buffers.foreach(_.close()) + allocator.close() + } + } + + /** Fills `n` rows of `schema` as Spark's readers lay them out, and checks them as above. */ + private def assertFilledRowsMatch( + schema: StructType, + n: Int, + nullFraction: Double, + offHeap: Boolean, + maxLength: Int = 6): Unit = { + val rnd = new Random(schema.hashCode) + val vectors = schema.fields.map(f => newVector(n, f.dataType, offHeap = false)) + try { + schema.fields.zip(vectors).foreach { case (field, v) => + fill(v, field.dataType, n, rnd, nullFraction, reversed = false, maxLength) + } + val batch = new ColumnarBatch(vectors.toArray[ColumnVector], n) + assertUnsafeRowsMatchGeneric(n, schema, offHeap)(batch.getRow) + } finally { + vectors.foreach(_.close()) + } + } + + for (offHeap <- Seq(false, true); nullFraction <- Seq(0.0, 0.2, 1.0)) { + test(s"unsafe rows match the generic row path: offHeap=$offHeap, nulls=$nullFraction") { + val types = primitiveTypes ++ nestedTypes ++ moreNestedTypes + val schemas = types.map(t => new StructType().add("c", t)) :+ + // A row wider than one word of null bits, with every type in it. + StructType(types.zipWithIndex.flatMap { case (t, i) => + Seq(StructField(s"a$i", t), StructField(s"b$i", t), StructField(s"c$i", t)) + }) + schemas.foreach { schema => + withClue(s"${schema.simpleString}: ") { + assertFilledRowsMatch(schema, numRows, nullFraction, offHeap) + } + } + } + } + + test("unsafe rows match the generic row path with collections past 64 elements") { + // An unsafe array keeps its null bits in 64-bit words, so long collections span several, + // and their elements land at every bit offset of the Arrow validity buffer. + val types = (nestedTypes ++ moreNestedTypes).collect { case t @ (_: ArrayType | _: MapType) => + t + } + for (dataType <- types; nullFraction <- Seq(0.0, 0.3)) { + withClue(s"$dataType nulls=$nullFraction: ") { + val schema = new StructType().add("c", dataType) + assertFilledRowsMatch(schema, 24, nullFraction, offHeap = false, maxLength = 150) + } + } + } + + /** + * Whether `target` is reachable from `root` through the fields of the Arrow writers and of any + * Spark unsafe views they hold. Arrow's vectors hold only Arrow memory, so they are skipped. + */ + private def reachable(root: AnyRef, target: AnyRef): Boolean = { + val writerPackage = classOf[ArrowWriter].getPackage.getName + "." + val seen = Collections.newSetFromMap(new IdentityHashMap[AnyRef, java.lang.Boolean]) + val pending = new ArrayDeque[AnyRef] + def pushFields(o: AnyRef): Unit = { + var c: Class[_] = o.getClass + while (c != null) { + c.getDeclaredFields.foreach { field => + if (!field.getType.isPrimitive && !Modifier.isStatic(field.getModifiers)) { + field.setAccessible(true) + val value = field.get(o) + if (value != null) { + pending.push(value) + } + } + } + c = c.getSuperclass + } + } + pending.push(root) + while (!pending.isEmpty) { + val o = pending.pop() + if (o eq target) { + return true + } + if (seen.add(o)) { + o match { + case objects: Array[AnyRef] => objects.foreach(x => if (x != null) pending.push(x)) + case _: UnsafeArrayData | _: UnsafeMapData | _: UnsafeRow => pushFields(o) + case _ if o.getClass.getName.startsWith(writerPackage) => pushFields(o) + case _ => + } + } + } + false + } + + test("the row path keeps no reference to an unsafe row once it is written") { + // Arrays, maps and structs at the top level, inside one another and as array elements. A view + // reused across values would keep the last row it read, and the row's memory, reachable. + val schema = new StructType() + .add("a", ArrayType(ArrayType(IntegerType))) + .add("s", new StructType().add("x", StringType).add("y", ArrayType(LongType))) + .add("m", MapType(StringType, new StructType().add("z", BinaryType))) + .add("as", ArrayType(new StructType().add("w", IntegerType))) + .add("am", ArrayType(MapType(IntegerType, StringType))) + val n = 20 + val rnd = new Random(7) + val vectors = schema.fields.map(f => newVector(n, f.dataType, offHeap = false)) + val allocator = new RootAllocator(Long.MaxValue) + val root = VectorSchemaRoot.create(Utils.toArrowSchema(schema, "UTC"), allocator) + try { + schema.fields.zip(vectors).foreach { case (field, v) => + fill(v, field.dataType, n, rnd, nullFraction = 0.0, reversed = false) + } + val batch = new ColumnarBatch(vectors.toArray[ColumnVector], n) + val project = UnsafeProjection.create(schema) + val writer = ArrowWriter.create(root, n) + (0 until n).foreach { i => + val row = project(batch.getRow(i)).copy() + writer.write(row) + withClue(s"row $i: ") { + reachable(writer, row.getBaseObject) shouldBe false + } + } + } finally { + root.close() + allocator.close() + vectors.foreach(_.close()) + } + } + + test("unsafe rows copy strings and binaries of every length") { + // Short values are copied a word at a time and long ones in one call, so this covers both + // sides of the threshold and every remainder. + val schema = new StructType() + .add("s", StringType) + .add("b", BinaryType) + .add("a", ArrayType(StringType)) + .add("m", MapType(StringType, BinaryType)) + val rows = (0 to 140).map { length => + if (length % 17 == 5) { + new GenericInternalRow(4) + } else { + val bytes = Array.tabulate[Byte](length)(i => (i * 31 + length).toByte) + val text = UTF8String.fromBytes(bytes) + val half = UTF8String.fromBytes(bytes, length / 3, length / 2) + new GenericInternalRow( + Array[Any]( + text, + bytes, + new GenericArrayData(Array[Any](text, null, half)), + new ArrayBasedMapData( + new GenericArrayData(Array[Any](text, UTF8String.fromString(s"key-$length"))), + new GenericArrayData(Array[Any](null, bytes))))) + } + } + Seq(false, true).foreach { offHeap => + assertUnsafeRowsMatchGeneric(rows.size, schema, offHeap)(rows) + } + } } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometStringWriterSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometStringWriterSuite.scala index 1e4230fb199..3f23851d76f 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometStringWriterSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometStringWriterSuite.scala @@ -22,6 +22,7 @@ package org.apache.spark.sql.comet.execution.arrow import java.nio.charset.StandardCharsets.UTF_8 import java.nio.file.{Files, Paths} +import scala.reflect.ClassTag import scala.util.Using import org.scalatest.funsuite.AnyFunSuite @@ -32,9 +33,10 @@ import org.apache.arrow.memory.{ArrowBuf, OutOfMemoryException, RootAllocator} import org.apache.arrow.vector.{VarCharVector, VectorSchemaRoot} import org.apache.arrow.vector.util.OversizedAllocationException import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.GenericInternalRow +import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeArrayData, UnsafeProjection, UnsafeRow} +import org.apache.spark.sql.catalyst.util.GenericArrayData import org.apache.spark.sql.comet.util.Utils -import org.apache.spark.sql.types.{StringType, StructField, StructType} +import org.apache.spark.sql.types.{ArrayType, StringType, StructField, StructType} import org.apache.spark.unsafe.types.UTF8String class CometStringWriterSuite extends AnyFunSuite with Matchers { @@ -299,4 +301,85 @@ class CometStringWriterSuite extends AnyFunSuite with Matchers { allocator.getAllocatedMemory shouldBe 0L } } + + private def unsafeRow(value: UTF8String): UnsafeRow = + UnsafeProjection.create(schema)(row(value)).copy() + + /** An unsafe array of `values`, in a row of its own. */ + private def unsafeArray(values: UTF8String*): UnsafeArrayData = { + val arraySchema = StructType(Seq(StructField("texts", ArrayType(StringType)))) + UnsafeProjection + .create(arraySchema)( + new GenericInternalRow(Array[Any](new GenericArrayData(values.toArray[Any])))) + .copy() + .getArray(0) + } + + /** `offsetAndSize`, as an unsafe row or array packs it, with the size made negative. */ + private def negativeSize(offsetAndSize: Long): Long = + (offsetAndSize & ~0xffffffffL) | 0xffffffffL + + /** + * Runs `write` from row `index(vector)` of a vector holding one value at row 0, whose offsets + * then end at `end`, and checks that it throws `E` having changed and allocated nothing. + */ + private def assertRejected[E <: Exception: ClassTag](end: Int, index: VarCharVector => Int)( + write: StringWriter => Unit): E = { + Using.resource(new RootAllocator(1024 * 1024)) { allocator => + val rejected = Using.Manager { use => + val vector = use(new VarCharVector("text", allocator)) + vector.allocateNew(8, 4) + vector.setSafe(0, Array[Byte](42)) + vector.getOffsetBuffer.setInt(4L, end) + val writer = new StringWriter(vector) + val at = index(vector) + writer.count = at + val allocated = allocator.getAllocatedMemory + val e = intercept[E](write(writer)) + allocator.getAllocatedMemory shouldBe allocated + writer.count shouldBe at + vector.getLastSet shouldBe 0 + vector.getOffsetBuffer.getInt(4L) shouldBe end + vector.getOffsetBuffer.getInt(8L) shouldBe 0 + vector.getDataBuffer.getByte(0L) shouldBe 42.toByte + if (at < vector.getValueCapacity) vector.isNull(at) shouldBe true + e + }.get + allocator.getAllocatedMemory shouldBe 0L + rejected + } + } + + test("unsafe row string offset overflow throws before anything changes or is allocated") { + // Unlike Arrow's setters, the unsafe paths check before growing the offsets. + val nearlyFull = Int.MaxValue - 5 + val value = unsafeRow(UTF8String.fromBytes(Array[Byte](1, 2, 3, 4, 5, 6))) + Seq[VarCharVector => Int](_ => 1, _.getValueCapacity).foreach { index => + assertRejected[OversizedAllocationException](nearlyFull, index) { + _.writeUnsafeRowField(value, 0) + } + } + // An array's elements are all checked before any of them is written. + val fitsThenOverflows = unsafeArray( + UTF8String.fromBytes(Array[Byte](1, 2)), + null, + UTF8String.fromBytes(Array[Byte](3, 4, 5, 6))) + Seq[VarCharVector => Int](_ => 1, _.getValueCapacity - 1).foreach { index => + assertRejected[OversizedAllocationException](nearlyFull, index) { + _.writeArrayElements(fitsThenOverflows) + } + } + } + + test("unsafe row negative string length is rejected before reserving or copying") { + val value = unsafeRow(UTF8String.fromString("abc")) + value.setLong(0, negativeSize(value.getLong(0))) + val array = unsafeArray(UTF8String.fromString("ab"), UTF8String.fromString("cd")) + array.setLong(1, negativeSize(array.getLong(1))) + Seq[StringWriter => Unit](_.writeUnsafeRowField(value, 0), _.writeArrayElements(array)) + .foreach { write => + assertRejected[IllegalArgumentException](1, _ => 1)(write).getMessage should + include("String length must be non-negative") + } + } } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala index a63560f3376..1d217bf6a25 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala @@ -265,6 +265,26 @@ class CometCelebornShufflePlanningSuite extends CometTestBase { } } + test("the shuffle input conversion leaves a Spark child to Spark's shuffle") { + // The conversion only takes over from the JVM columnar shuffle, which Celeborn does not + // have, and Celeborn's native shuffle needs a native child. Without Celeborn, auto and jvm + // would convert this child. Native mode has no JVM columnar shuffle either, so the reason + // is what shows that Celeborn kept Spark's shuffle there. + Seq("auto", "jvm", "native").foreach { mode => + withSQLConf( + CometConf.COMET_SHUFFLE_MODE.key -> mode, + CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.key -> "true") { + val leaf = sparkLeaf() + val exchange = ShuffleExchangeExec(HashPartitioning(leaf.output, 2), leaf) + assert(!CometShuffleExchangeExec.convertsInputForNativeShuffle(exchange), mode) + assert(reasons(exchange).exists(_.contains("Celeborn")), s"$mode: ${reasons(exchange)}") + val planned = CometExecRule(spark)(exchange) + assert(planned.isInstanceOf[ShuffleExchangeExec], planned) + assert(collect(planned) { case c: CometSparkToColumnarExec => c }.isEmpty, planned) + } + } + } + test("fallback is sticky after AQE reshapes the child or native mode becomes available") { val child = nativeChild() val exchange = ShuffleExchangeExec(SinglePartition, child) diff --git a/spark/src/test/spark-4.1/org/apache/spark/sql/CometCollationSuite.scala b/spark/src/test/spark-4.1+/org/apache/spark/sql/CometCollationSuite.scala similarity index 100% rename from spark/src/test/spark-4.1/org/apache/spark/sql/CometCollationSuite.scala rename to spark/src/test/spark-4.1+/org/apache/spark/sql/CometCollationSuite.scala diff --git a/spark/src/test/spark-4.1/org/apache/spark/sql/benchmark/CometTimeExtractBenchmark.scala b/spark/src/test/spark-4.1/org/apache/spark/sql/benchmark/CometTimeExtractBenchmark.scala index 025fe43d459..38d4ac8b549 100644 --- a/spark/src/test/spark-4.1/org/apache/spark/sql/benchmark/CometTimeExtractBenchmark.scala +++ b/spark/src/test/spark-4.1/org/apache/spark/sql/benchmark/CometTimeExtractBenchmark.scala @@ -85,6 +85,7 @@ object CometTimeExtractBenchmark extends CometBenchmarkBase { private def runKernels(precision: Int, nulls: Boolean): Unit = { val size = 10000 val time = new TimeNanoVector("t", CometArrowAllocator) + val digest = new VarBinaryVector("digest", CometArrowAllocator) val serialized = new VarBinaryVector("expr", CometArrowAllocator) try { time.allocateNew(size) @@ -114,11 +115,14 @@ object CometTimeExtractBenchmark extends CometBenchmarkBase { val buffer = SparkEnv.get.closureSerializer.newInstance().serialize(expr) val bytes = new Array[Byte](buffer.remaining()) buffer.get(bytes) + digest.allocateNew() + digest.setSafe(0, CometScalaUDFCodegen.digest(bytes)) + digest.setValueCount(1) serialized.allocateNew() serialized.setSafe(0, bytes) serialized.setValueCount(1) val dispatcher = new CometScalaUDFCodegen() - val inputs = Array[org.apache.arrow.vector.ValueVector](serialized, time) + val inputs = Array[org.apache.arrow.vector.ValueVector](digest, serialized, time) val result = dispatcher.evaluate(inputs, size) try { val vector = @@ -170,6 +174,7 @@ object CometTimeExtractBenchmark extends CometBenchmarkBase { } } } finally { + digest.close() serialized.close() time.close() } diff --git a/spark/src/test/spark-4.2/org/apache/comet/iceberg/IcebergTransactionalCommitTestShim.scala b/spark/src/test/spark-4.2/org/apache/comet/iceberg/IcebergTransactionalCommitTestShim.scala index adb0b05f8c5..64d78504933 100644 --- a/spark/src/test/spark-4.2/org/apache/comet/iceberg/IcebergTransactionalCommitTestShim.scala +++ b/spark/src/test/spark-4.2/org/apache/comet/iceberg/IcebergTransactionalCommitTestShim.scala @@ -42,7 +42,8 @@ object IcebergTransactionalCommitTestShim { null.asInstanceOf[BatchWrite], null.asInstanceOf[Write], () => (), - LocalTableScanExec(Nil, Nil, None)) + LocalTableScanExec(Nil, Nil, None), + command = Some(DeltaUpdate)) val original = IcebergCommitPlanShim .wrap(commit) .asInstanceOf[IcebergTransactionalCommitExec] @@ -58,6 +59,9 @@ object IcebergTransactionalCommitTestShim { val attached = transformed.asInstanceOf[IcebergTransactionalCommitExec] assert(attached.transaction.contains(transaction)) - assert(attached.clone().transaction.contains(transaction)) + assert(attached.command.contains(DeltaUpdate)) + val cloned = attached.clone() + assert(cloned.transaction.contains(transaction)) + assert(cloned.command.contains(DeltaUpdate)) } } diff --git a/spark/src/test/spark-4.x/org/apache/comet/parquet/CometEmptyRelationParquetWriterSuite.scala b/spark/src/test/spark-4.x/org/apache/comet/parquet/CometEmptyRelationParquetWriterSuite.scala index 04585612ab0..c739e420315 100644 --- a/spark/src/test/spark-4.x/org/apache/comet/parquet/CometEmptyRelationParquetWriterSuite.scala +++ b/spark/src/test/spark-4.x/org/apache/comet/parquet/CometEmptyRelationParquetWriterSuite.scala @@ -64,12 +64,8 @@ class CometEmptyRelationParquetWriterSuite extends CometParquetWriterTestBase { } if (nativeEmpty) { assert(collect(plan) { case e: CometEmptyRelationExec => e }.nonEmpty) - // A native empty relation is a zero-partition RDD, which is what - // CometDataWritingCommand declines on the Spark 3.x path: that writer only maps - // existing partitions, so no task runs and no file is written. Spark 4.0+ goes - // through the WriteFilesExec seam instead, where CometWriteFilesExec swaps in a - // dummy single-partition RDD exactly as Spark's own WriteFilesExec does, so - // partition 0 still writes the schema-only file the readback above needs. + // CometWriteFilesExec replaces the zero-partition native empty relation + // with one empty partition, so a task writes the schema-only file above. assertHasCometNativeWriteExec(plan) } else { // Spark's own EmptyRelationExec is not a Comet operator, so