diff --git a/.github/workflows/README.md b/.github/workflows/README.md index 7e2436df554..1976ca2632a 100644 --- a/.github/workflows/README.md +++ b/.github/workflows/README.md @@ -24,9 +24,9 @@ ruleset in `.asf.yaml`. That splits CI into three tiers: and one Iceberg version, both the default profile's. - **Nightly tier** (`nightly`): the regression sweep of everything else, once a day against `main` as it stands. The Comet test suites against the other - four Spark profiles, Spark SQL on Spark 3.5, 4.0 and 4.2, and Iceberg - 1.8/1.9/1.10. See [Nightly tier](#nightly-tier) below for how a failure - surfaces. + four Spark profiles, Spark SQL on Spark 3.5, 4.0 and 4.2, Iceberg + 1.8/1.9/1.10, and the Delta contrib suites on Spark 3.5, 4.0 and 4.1. See + [Nightly tier](#nightly-tier) below for how a failure surfaces. Every queue-only and nightly job has a `run-*` label that opts a pull request into it early, listed in the diagram below. The Lint Java matrix compiles @@ -129,6 +129,9 @@ tiers partition the list and that the `pr` tier is exactly the default profile. iceberg_1_8 run-iceberg-tests iceberg_1_9 run-iceberg-tests iceberg_1_10 run-iceberg-tests + delta_3_5 run-delta-tests + delta_4_0 run-delta-tests + delta_4_1 run-delta-tests | | | +-----------------------------------+-----------------------------------+ @@ -165,6 +168,9 @@ tiers partition the list and that the `pr` tier is exactly the default profile. | `pr_benchmark_check` | merge group, **or** PR with `run-benchmark-check` | benchmark sources only | | `delta_build_gate` | merge group, **or** PR with `run-delta-build-gate` | main sources, poms, `contrib/delta` | | `pyarrow_udf_test` | merge group, **or** PR with `run-pyarrow-udf-tests` | map-in-batch and Python runner code | +| `delta_3_5` | nightly, **or** PR with `run-delta-tests` | Delta contrib and native sources | +| `delta_4_0` | nightly, **or** PR with `run-delta-tests` | Delta contrib and native sources | +| `delta_4_1` | nightly, **or** PR with `run-delta-tests` | Delta contrib and native sources | | `docs` | push to main, paths matched, **or** dispatch on `main` | `.asf.yaml`, `docs/**`, `docs.yaml` | | `spark_3_5` | nightly, **or** PR with `run-spark-3.5-tests` | Spark 3.5 sources | | `spark_4_1` | merge group, **or** PR with `run-spark-4.1-tests`; the `sql_hive` shards alone with `run-spark-4.1-hive-tests` | Spark 4.1 sources | @@ -303,6 +309,7 @@ umbrella doesn't watch, or operate independently of the rest of CI: | `pr_build_macos.yml` | `pr_build_macos` | | `pr_benchmark_check.yml` | `pr_benchmark_check` | | `delta_build_gate.yml` | `delta_build_gate` | +| `delta_contrib_test.yml` | `delta_3_5`, `delta_4_0`, `delta_4_1` | | `pyarrow_udf_test.yml` | `pyarrow_udf_test` | | `docs.yaml` | `docs` | | `spark_sql_test_reusable.yml` | `spark_3_4`, `spark_3_5`, `spark_4_0`, `spark_4_1`, `spark_4_2` | diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 075dd292366..936b804272d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -263,6 +263,9 @@ jobs: iceberg_1_9: ${{ steps.compute.outputs.iceberg_1_9 }} iceberg_1_10: ${{ steps.compute.outputs.iceberg_1_10 }} iceberg_1_11: ${{ steps.compute.outputs.iceberg_1_11 }} + delta_3_5: ${{ steps.compute.outputs.delta_3_5 }} + delta_4_0: ${{ steps.compute.outputs.delta_4_0 }} + delta_4_1: ${{ steps.compute.outputs.delta_4_1 }} steps: - uses: actions/checkout@v7 with: @@ -469,6 +472,43 @@ jobs: spark-full: '3.5.9' java: 17 + delta_3_5: + name: Delta Contrib Tests (Spark 3.5) + needs: changes + permissions: + contents: read + # Nightly by default; PRs need the `run-delta-tests` label. + if: needs.changes.outputs.delta_3_5 == 'true' + uses: ./.github/workflows/delta_contrib_test.yml + with: + spark: '3.5' + java: '17' + + delta_4_0: + name: Delta Contrib Tests (Spark 4.0) + needs: changes + permissions: + contents: read + # Nightly by default; PRs need the `run-delta-tests` label. + if: needs.changes.outputs.delta_4_0 == 'true' + uses: ./.github/workflows/delta_contrib_test.yml + with: + spark: '4.0' + java: '17' + + delta_4_1: + name: Delta Contrib Tests (Spark 4.1) + needs: changes + permissions: + contents: read + # Nightly by default; PRs need the `run-delta-tests` label. No Spark 4.2 + # caller: the contrib profile is dormant until a Delta release supports it. + if: needs.changes.outputs.delta_4_1 == 'true' + uses: ./.github/workflows/delta_contrib_test.yml + with: + spark: '4.1' + java: '17' + spark_4_0: name: Spark SQL Tests (Spark 4.0) needs: changes @@ -616,6 +656,9 @@ jobs: - pyarrow_udf_test - spark_3_4 - spark_3_5 + - delta_3_5 + - delta_4_0 + - delta_4_1 - spark_4_0 - spark_4_1 - spark_4_2 @@ -684,7 +727,7 @@ jobs: failed.length ? 'Jobs that did not succeed:' : 'No individual job reported a failure; see the run for details.', ...failed, '', - 'The nightly tier runs the Spark and Iceberg versions the merge queue does not', + 'The nightly tier runs the Spark, Iceberg and Delta versions the merge queue does not', '(see `POLICY` in `dev/ci/compute-changes.py`). Close this issue once the failure', 'is understood; the next red nightly opens a new one.', ].join('\n'); diff --git a/.github/workflows/delta_build_gate.yml b/.github/workflows/delta_build_gate.yml index cd7fd2ac19c..caf8a37ae3f 100644 --- a/.github/workflows/delta_build_gate.yml +++ b/.github/workflows/delta_build_gate.yml @@ -15,12 +15,15 @@ # specific language governing permissions and limitations # under the License. -# Build-gate verification for the optional contrib/delta integration. +# Build-gate verification for the Delta split of the native library and the JVM build. # # A single cheap job that runs dev/verify-contrib-delta-gate.sh, which proves: -# - the DEFAULT cargo / mvn / dylib build carries ZERO Delta surface -# (no comet-contrib-delta / delta_kernel in the cargo tree, no io.delta in -# the effective pom, no Delta symbols in libcomet); and +# - the kernel-backed `comet-contrib-delta` crate stays out of every shipped +# build (no comet-contrib-delta / delta_kernel in the default or +# --no-default-features cargo tree, no io.delta in the effective pom, no +# contrib symbols in libcomet), while the small default-on `delta` feature +# (deletion-vector decoding) is deliberately in and its symbol footprint is +# pinned; and # - the gated build (`--features contrib-delta` / `-Pcontrib-delta`) compiles # and pulls the right delta-spark per Spark profile. # diff --git a/.github/workflows/delta_contrib_test.yml b/.github/workflows/delta_contrib_test.yml new file mode 100644 index 00000000000..8bf1d14b411 --- /dev/null +++ b/.github/workflows/delta_contrib_test.yml @@ -0,0 +1,270 @@ +# 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. + +name: Delta Contrib Tests + +# Reusable: invoked by ci.yml once per Spark profile. Triggering, path +# filters, and concurrency live in the umbrella workflow. +on: + workflow_call: + inputs: + spark: + description: 'Spark minor version, e.g. 3.5' + required: true + type: string + java: + description: 'JDK major version, e.g. 17' + required: true + type: string + +permissions: + contents: read + +env: + RUST_VERSION: stable + RUST_BACKTRACE: 1 + # Force GNU ld on Linux: rust-lld cannot resolve -ljvm against the Zulu JDK + # layout installed by setup-java (same rationale as pr_build_linux.yml). + RUSTFLAGS: "-Clink-arg=-fuse-ld=bfd" + # The container's default locale is POSIX, which makes the JVM's file-path encoder + # reject non-ASCII partition directory names the suites create. + LANG: "C.UTF-8" + LC_ALL: "C.UTF-8" + +jobs: + + contrib-delta: + name: Delta contrib (Spark ${{ inputs.spark }}) + runs-on: ubuntu-24.04 + container: + image: amd64/rust + steps: + - uses: actions/checkout@v7 + + - name: Setup Rust & Java toolchain + uses: ./.github/actions/setup-builder + with: + rust-version: ${{ env.RUST_VERSION }} + jdk-version: ${{ inputs.java }} + + - name: Restore Maven dependencies + id: maven-cache + uses: actions/cache/restore@v6 + with: + path: | + ~/.m2/repository + /root/.m2/repository + key: ${{ runner.os }}-java-maven-${{ hashFiles('**/pom.xml') }}-delta-${{ inputs.spark }} + restore-keys: | + ${{ runner.os }}-java-maven- + + - name: Restore Cargo cache + uses: actions/cache/restore@v6 + with: + path: | + ~/.cargo/registry + ~/.cargo/git + native/target + key: ${{ runner.os }}-cargo-ci-${{ hashFiles('native/**/Cargo.lock', 'native/**/Cargo.toml') }}-${{ hashFiles('native/**/*.rs') }} + restore-keys: | + ${{ runner.os }}-cargo-ci-${{ hashFiles('native/**/Cargo.lock', 'native/**/Cargo.toml') }}- + + - name: Build native library (CI profile) + run: | + cd native + # `delta` is in the default feature set, so no feature flag is needed. + cargo build --profile ci + env: + # Must match the flags spark_sql_test_reusable.yml builds with: + # cargo folds RUSTFLAGS into its fingerprints, so any divergence + # would make the shared cargo cache restore without ever hitting. + RUSTFLAGS: "-Ctarget-cpu=x86-64-v3 -Clink-arg=-fuse-ld=bfd" + + # No cache save: this workflow never runs on a push to main, so the restore above + # falls back to the cache build_linux saves there. + + - name: Stage native library at release path + run: | + # Maven's -Prelease profile (activated below) expects libcomet.so + # under native/target/release/; --profile ci builds it under + # native/target/ci/ instead (same as the other native-building + # workflows), so copy it into place. + mkdir -p native/target/release + cp native/target/ci/libcomet.so native/target/release/libcomet.so + + - name: Install Comet core jars + run: | + # The delta profile is what produces the spark test-jar the contrib suites depend on. + ./mvnw -B -q -Prelease -Pspark-${{ inputs.spark }},delta install -pl common,spark -DskipTests -Dspotless.check.skip=true + + - name: Run Delta contrib test suites + run: | + SPARK_HOME=$(pwd) COMET_CONF_DIR=$(pwd)/conf ./mvnw -B -Prelease -Pspark-${{ inputs.spark }},delta test -pl contrib/delta-spark + + # Saved only on main: an entry written from a pull request or a + # `gh-readonly-queue/*` branch cannot be restored by any later run, and + # it evicts main's from the shared budget. See "Large caches are written + # on main only" in README.md. + - name: Save Maven dependencies + if: ${{ github.ref == 'refs/heads/main' && steps.maven-cache.outputs.cache-hit != 'true' }} + uses: actions/cache/save@v6 + with: + path: | + ~/.m2/repository + /root/.m2/repository + key: ${{ runner.os }}-java-maven-${{ hashFiles('**/pom.xml') }}-delta-${{ inputs.spark }} + + # The container job above has no Docker socket, so CometDeltaS3Suite cancels itself + # there. This job runs it on the plain runner, where Testcontainers can start MinIO. + contrib-delta-s3: + name: Delta contrib S3 suite (Spark 3.5, MinIO) + # Pinned to Spark 3.5, so only that caller runs it: once per ci.yml run. + if: inputs.spark == '3.5' + runs-on: ubuntu-24.04 + steps: + - uses: actions/checkout@v7 + + # setup-builder runs apt-get as root inside the container jobs; the hosted runner needs + # sudo for that, so the same steps are spelled out here. + - name: Install build dependencies + run: | + sudo apt-get update + sudo apt-get install -y protobuf-compiler clang + + - name: Install JDK 17 + uses: actions/setup-java@v6 + with: + distribution: 'zulu' + java-version: "17" + + - name: Setup Rust toolchain + run: | + rustup toolchain install ${{ env.RUST_VERSION }} + rustup default ${{ env.RUST_VERSION }} + rustup component add rustfmt clippy + + - name: Restore Maven dependencies + id: maven-cache + uses: actions/cache/restore@v6 + with: + path: | + ~/.m2/repository + key: ${{ runner.os }}-java-maven-${{ hashFiles('**/pom.xml') }}-delta-3.5 + restore-keys: | + ${{ runner.os }}-java-maven- + + - name: Restore Cargo cache + uses: actions/cache/restore@v6 + with: + path: | + ~/.cargo/registry + ~/.cargo/git + native/target + key: ${{ runner.os }}-cargo-ci-${{ hashFiles('native/**/Cargo.lock', 'native/**/Cargo.toml') }}-${{ hashFiles('native/**/*.rs') }} + restore-keys: | + ${{ runner.os }}-cargo-ci-${{ hashFiles('native/**/Cargo.lock', 'native/**/Cargo.toml') }}- + + - name: Build native library (CI profile) + run: | + cd native + cargo build --profile ci + env: + RUSTFLAGS: "-Ctarget-cpu=x86-64-v3 -Clink-arg=-fuse-ld=bfd" + + - name: Stage native library at release path + run: | + mkdir -p native/target/release + cp native/target/ci/libcomet.so native/target/release/libcomet.so + + - name: Install Comet core jars + run: | + ./mvnw -B -q -Prelease -Pspark-3.5,delta install -pl common,spark -DskipTests -Dspotless.check.skip=true + + - name: Run the MinIO-backed S3 suite + run: | + docker info > /dev/null + SPARK_HOME=$(pwd) COMET_CONF_DIR=$(pwd)/conf ./mvnw -B -Prelease -Pspark-3.5,delta test -pl contrib/delta-spark -DwildcardSuites=org.apache.comet.contrib.delta.CometDeltaS3Suite + env: + # This job exists only to run the suite, so a Docker or MinIO failure must fail + # it instead of cancelling every test, which the container job above relies on. + COMET_DELTA_S3_REQUIRED: "1" + + # Saved only on main, for the same reason as the container job above. + - name: Save Maven dependencies + if: ${{ github.ref == 'refs/heads/main' && steps.maven-cache.outputs.cache-hit != 'true' }} + uses: actions/cache/save@v6 + with: + path: | + ~/.m2/repository + key: ${{ runner.os }}-java-maven-${{ hashFiles('**/pom.xml') }}-delta-3.5 + + # `delta` is in the default feature set, so no regular job builds without it; + # this keeps the feature-off build and its "built without the delta feature" + # error arm (planner.rs cfg(not(feature = "delta"))) from becoming dead code. + feature-off-build: + name: Feature-off native build + # Spark-independent, so only the 3.5 caller runs it: once per ci.yml run. + if: inputs.spark == '3.5' + runs-on: ubuntu-24.04 + container: + image: amd64/rust + steps: + - uses: actions/checkout@v7 + + - name: Setup Rust & Java toolchain + uses: ./.github/actions/setup-builder + with: + rust-version: ${{ env.RUST_VERSION }} + jdk-version: "17" + + - name: Restore Cargo cache + uses: actions/cache/restore@v6 + with: + path: | + ~/.cargo/registry + ~/.cargo/git + native/target + key: ${{ runner.os }}-cargo-ci-${{ hashFiles('native/**/Cargo.lock', 'native/**/Cargo.toml') }}-${{ hashFiles('native/**/*.rs') }} + restore-keys: | + ${{ runner.os }}-cargo-ci-${{ hashFiles('native/**/Cargo.lock', 'native/**/Cargo.toml') }}- + + - name: Test the delta-off error path + run: | + cd native + # The test binary links JNI; libjvm.so must be resolvable at load + # time (same as .github/actions/rust-test). + export LD_LIBRARY_PATH=${JAVA_HOME}/lib/server:${LD_LIBRARY_PATH} + cargo test -p datafusion-comet --no-default-features --features hdfs-opendal delta_scan + + # The contrib's dev tooling (benchmark and regression-harness scripts) is + # Python; keep it import-clean on one supported interpreter. + dev-scripts-python: + name: Delta dev scripts (Python) + # Spark-independent, so only the 3.5 caller runs it: once per ci.yml run. + if: inputs.spark == '3.5' + runs-on: ubuntu-24.04 + steps: + - uses: actions/checkout@v7 + + - name: Setup Python + uses: actions/setup-python@v6 + with: + python-version: "3.12" + + - name: Byte-compile contrib dev scripts + run: | + python -m compileall -q contrib/delta-spark/dev diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 192dc8ffb32..499ad824bc5 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -242,8 +242,10 @@ jobs: # Compiles main and test sources with -Pstrict-warnings, which promotes scalac warnings to # errors, so a change that reintroduces a cleared warning (an Int widened into a Long metric, - # a discarded builder result) fails here rather than accumulating. Spark 3.5 only: the profile - # passes on the Scala 2.12 profiles, and the Scala 2.13 remainder is tracked in + # a discarded builder result) fails here rather than accumulating. It also compiles the Delta + # contrib (-Pdelta) against Spark 3.5, so a change that breaks the contrib build on Spark 3.5 + # fails on its own pull request rather than in the nightly Delta suites. Spark 3.5 only: the + # profile passes on the Scala 2.12 profiles, and the Scala 2.13 remainder is tracked in # https://github.com/apache/datafusion-comet/issues/5893. strict-scala-warnings: needs: lint @@ -273,7 +275,7 @@ jobs: ${{ runner.os }}-java-maven- - name: Compile with strict warnings - run: ./mvnw -B test-compile -Pspark-3.5 -Pstrict-warnings -DskipTests + run: ./mvnw -B test-compile -Pspark-3.5,delta -Pstrict-warnings -DskipTests # Saved only on main: an entry written from a pull request or a # `gh-readonly-queue/*` branch cannot be restored by any later run, and diff --git a/contrib/delta-spark/README.md b/contrib/delta-spark/README.md new file mode 100644 index 00000000000..89664d3d61a --- /dev/null +++ b/contrib/delta-spark/README.md @@ -0,0 +1,70 @@ + + +# Comet Delta Lake Contrib (experimental) + +Native Delta Lake reads for Comet. Delta tables are scanned through Comet's +existing native Parquet reader, so they get row-group pruning, page-index +pruning, and filter pushdown, with deletion vectors applied inside the scan. + +Support is experimental and explicitly opt-in. Two things are required: + +1. This module's jar (`comet-contrib-delta-spark`) on the classpath, alongside + `delta-spark`. It is never bundled into `comet-spark`; without it, Comet + has no Delta surface at all. +2. `spark.comet.scan.delta.enabled=true`. The default is `false`, so the jar + alone does nothing. + +Unsupported tables and features fall back to Spark's reader. See the +[user guide](https://datafusion.apache.org/comet/user-guide/latest/delta.html) +for configuration details. + +## Supported versions + +| Spark | Delta | Status | +| ----- | -------------- | --------------------------------------------- | +| 3.5 | 3.3.x | supported | +| 4.0 | 4.0.x | supported | +| 4.1 | 4.3.x | supported | +| 3.4 | delta-core 2.4 | not supported (older Delta, would need shims) | +| 4.2 | none released | inert until Delta ships a Spark 4.2 release | + +## Building and testing + +The module builds under the `delta` Maven profile. It resolves `comet-spark` +from the local Maven repository, so install `common` and `spark` from the same +checkout immediately before, as CI does, with the `delta` profile active so that +install produces the spark test-jar the contrib suites depend on; a stale sibling +install is the trap the contributor guide warns about: + +```shell +./mvnw -Pspark-3.5,delta install -pl common,spark -DskipTests +./mvnw -Pspark-3.5,delta install -pl contrib/delta-spark +``` + +Run the test suites the same way (`test` instead of `install` on the second +line). CI runs them via `.github/workflows/delta_contrib_test.yml`: nightly on +Spark 3.5, 4.0 and 4.1, or on a pull request carrying the `run-delta-tests` +label. +`CometDeltaS3Suite` starts MinIO through Testcontainers and cancels itself when +no Docker daemon is reachable; setting `COMET_DELTA_S3_REQUIRED=1`, as the +`contrib-delta-s3` CI job does, turns that cancel into a suite failure. + +`dev/` contains a benchmark script (`bench_delta_comet.py`) and a harness for +running Delta's own test suites against Comet (`run-delta-regression.sh`). diff --git a/contrib/delta-spark/dev/bench_delta_comet.py b/contrib/delta-spark/dev/bench_delta_comet.py new file mode 100644 index 00000000000..66003c39c3a --- /dev/null +++ b/contrib/delta-spark/dev/bench_delta_comet.py @@ -0,0 +1,272 @@ +# 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: page-level skipping on a DELTA table under three configurations. + + 1. stock -- plain Spark 3.5.6 + delta-spark 3.3.2 + 2. comet -- Comet enabled WITHOUT the Delta contrib (scan falls back to Spark) + 3. contrib -- Comet + comet-contrib-delta (native Delta scan) + +Writes a 20M-row table sorted by `ts` (4 files, zstd, small pages) as Delta, +optionally deletes a slice via DVs, then runs a 5%-wide range predicate and +reports the fraction of the table materialized by the scan plus wall time. + +Usage: python bench_delta_comet.py [--dv] [--subquery] + mode: stock | comet | contrib (jars/extensions injected by the wrapper script) + --subquery: bound the range predicate with scalar subqueries over a one-row + thresholds Delta table instead of literals. Same rows selected; exercises + the execution-time resolve-and-push path (which stock Spark 3.5 lacks: + FileSourceStrategy strips subquery predicates from scan dataFilters). +""" + +import os +import sys +import time + +from pyspark.sql import SparkSession +from pyspark.sql import functions as F + +ROWS = 20_000_000 +FILES = 4 +PRED_LO, PRED_HI = 0.475, 0.525 # 5% slice in the middle +# DV delete ranges: one nested inside the predicate slice, one far outside it. +DV_DELETE_LO, DV_DELETE_HI = 0.48, 0.49 + + +def build_session(mode: str) -> SparkSession: + extensions = "io.delta.sql.DeltaSparkSessionExtension" + if mode in ("comet", "contrib"): + extensions += ",org.apache.comet.CometSparkSessionExtensions" + b = ( + SparkSession.builder.appName(f"delta-comet-bench-{mode}") + .config("spark.sql.extensions", extensions) + .config("spark.sql.adaptive.enabled", "false") + .config( + "spark.sql.catalog.spark_catalog", + "org.apache.spark.sql.delta.catalog.DeltaCatalog", + ) + .config("spark.driver.memory", "6g") + .config("spark.sql.shuffle.partitions", "8") + .config("spark.ui.enabled", "false") + .config("spark.hadoop.parquet.page.size", str(64 * 1024)) + .config("spark.hadoop.parquet.block.size", str(32 * 1024 * 1024)) + ) + if mode in ("comet", "contrib"): + b = ( + b.config("spark.comet.enabled", "true") + .config("spark.comet.exec.enabled", "true") + .config("spark.comet.exec.shuffle.enabled", "true") + .config( + "spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager", + ) + .config("spark.memory.offHeap.enabled", "true") + .config("spark.memory.offHeap.size", "4g") + .config("spark.comet.explainFallback.enabled", "true") + ) + if mode == "contrib": + b = b.config("spark.comet.scan.delta.enabled", "true") + return b.getOrCreate() + + +def write_table(spark: SparkSession, path: str, with_dv: bool) -> None: + df = ( + spark.range(ROWS) + .withColumn("ts", F.col("id")) + .withColumn("payload", F.sha1(F.col("id").cast("string"))) + .repartitionByRange(FILES, "ts") + .sortWithinPartitions("ts") + ) + ( + df.write.format("delta") + .option("compression", "zstd") + .mode("overwrite") + .save(path) + ) + if with_dv: + spark.sql( + f"ALTER TABLE delta.`{path}` SET TBLPROPERTIES " + "('delta.enableDeletionVectors' = 'true')" + ) + lo = int(ROWS * DV_DELETE_LO) + hi = int(ROWS * DV_DELETE_HI) + spark.sql(f"DELETE FROM delta.`{path}` WHERE ts >= {lo} AND ts < {hi}") + + +def scan_metrics(plan): + """Walk the executed plan and pull metrics from the leaf scan node(s). + + Safe to call right after collect(): the Dataset caches its QueryExecution, + per-task SQLMetric accumulator updates are merged on the driver before the + job completes, and AQE is disabled so the executed plan is final. + """ + from py4j.protocol import Py4JError, Py4JJavaError + + out = {} + + def walk(node): + try: + name = node.nodeName() + if "Scan" in name: + metrics = node.metrics() + it = metrics.keysIterator() + while it.hasNext(): + k = it.next() + out.setdefault((name, k), metrics.get(k).get().value()) + for i in range(node.children().length()): + walk(node.children().apply(i)) + # innerChildren covers plan-in-plan nodes; entries may not be + # SparkPlans, so failures here are ignored rather than fatal. + inner = node.innerChildren() + for i in range(inner.length()): + walk(inner.apply(i)) + except (Py4JError, Py4JJavaError): + pass + + walk(plan) + return out + + +def has_native_scan_with_column(plan, column: str) -> bool: + """True if the executed plan (including subquery inner plans) contains a + CometDeltaNativeScan whose output includes `column`. Programmatic version of + the test suite's `output.exists(_.name == col)` check -- identifies the MAIN + table's scan by its distinctive column, since subquery mode adds trivial + thresholds-table scans that would fool any name-only or count-based check. + """ + from py4j.protocol import Py4JError, Py4JJavaError + + def walk(node) -> bool: + try: + if node.nodeName().startswith("CometDeltaNativeScan"): + attrs = node.output() + for i in range(attrs.length()): + if attrs.apply(i).name() == column: + return True + for i in range(node.children().length()): + if walk(node.children().apply(i)): + return True + inner = node.innerChildren() + for i in range(inner.length()): + if walk(inner.apply(i)): + return True + except (Py4JError, Py4JJavaError): + pass + return False + + return walk(plan) + + +def pred_bounds() -> tuple[int, int]: + """Single source of truth for the range bounds, so the literal and subquery + modes are guaranteed to select the same rows.""" + return int(ROWS * PRED_LO), int(ROWS * PRED_HI) + + +def write_thresholds(spark: SparkSession, thr_path: str) -> None: + lo, hi = pred_bounds() + spark.sql( + f"SELECT CAST({lo} AS BIGINT) AS lo, CAST({hi} AS BIGINT) AS hi" + ).write.format("delta").mode("overwrite").save(thr_path) + + +def run_query(spark: SparkSession, path: str, thr_path: str | None = None): + if thr_path is not None: + df = spark.sql( + f"SELECT count(*) AS n, sum(length(payload)) AS s FROM delta.`{path}` " + f"WHERE ts >= (SELECT lo FROM delta.`{thr_path}`) " + f"AND ts < (SELECT hi FROM delta.`{thr_path}`)" + ) + else: + lo, hi = pred_bounds() + df = ( + spark.read.format("delta") + .load(path) + .where((F.col("ts") >= lo) & (F.col("ts") < hi)) + .agg(F.count("*").alias("n"), F.sum(F.length("payload")).alias("s")) + ) + t0 = time.perf_counter() + row = df.collect()[0] + elapsed = time.perf_counter() - t0 + plan = df._jdf.queryExecution().executedPlan() + mets = scan_metrics(plan) + main_scan_native = has_native_scan_with_column(plan, "payload") + return row, elapsed, mets, plan.toString(), main_scan_native + + +def main(): + if len(sys.argv) < 3 or sys.argv[1] not in ("stock", "comet", "contrib"): + print(__doc__) + sys.exit(2) + mode, workdir = sys.argv[1], sys.argv[2] + with_dv = "--dv" in sys.argv + with_subquery = "--subquery" in sys.argv + path = f"{workdir}/delta_bench{'_dv' if with_dv else ''}" + thr_path = f"{workdir}/delta_bench_thr" if with_subquery else None + spark = build_session(mode) + spark.sparkContext.setLogLevel("WARN") + + if not os.path.exists(path + "/_delta_log"): + print(f"[bench] writing table to {path}") + write_table(spark, path, with_dv) + if thr_path is not None and not os.path.exists(thr_path + "/_delta_log"): + write_thresholds(spark, thr_path) + + try: + # warm-up then measured run + run_query(spark, path, thr_path) + row, elapsed, mets, plan_str, main_scan_native = run_query(spark, path, thr_path) + except BaseException: + spark.stop() + raise + + print(f"\n=== mode={mode} dv={with_dv} subquery={with_subquery} ===") + print(f"result: n={row['n']} sum={row['s']}") + print(f"wall_time_s: {elapsed:.3f}") + interesting = ( + "output_rows", + "numOutputRows", + "bytes_scanned", + "page_index_rows_pruned", + "page_index_rows_matched", + "row_groups_pruned_statistics", + "row_groups_matched_statistics", + "numFiles", + "filesSize", + ) + for (node, k), v in sorted(mets.items()): + if any(k == i for i in interesting): + print(f"metric: {node} :: {k} = {v}") + # rows materialized by the scan as fraction of table + scanned = [v for (n, k), v in mets.items() if k in ("output_rows", "numOutputRows")] + if scanned: + frac = max(scanned) / ROWS + print(f"scan_fraction: {frac:.4f}") + seen = {k for (_, k) in mets} + for key in ("output_rows", "numOutputRows"): + if key in seen: + break + else: + print("WARNING: no scan row metrics found; scan_fraction unavailable") + if mode == "contrib" and not main_scan_native: + print("WARNING: contrib mode but the main table's scan is not CometDeltaNativeScan!") + spark.stop() + + +if __name__ == "__main__": + main() diff --git a/contrib/delta-spark/dev/run-delta-regression.sh b/contrib/delta-spark/dev/run-delta-regression.sh new file mode 100755 index 00000000000..44cffb8d8ac --- /dev/null +++ b/contrib/delta-spark/dev/run-delta-regression.sh @@ -0,0 +1,176 @@ +#!/bin/bash +# +# 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. +# +# Run Delta Lake's own Spark test suites against a Comet build with the +# native Delta scan enabled. Clones delta at $DELTA_VERSION into $WORKDIR, +# injects Comet into the test SparkSession (DeltaSQLCommandTest) and the +# test classpath (unmanagedJars), then runs the given testOnly selectors. +# +# Usage: +# COMET_JARS=/path/comet-spark.jar,/path/comet-contrib-delta.jar,/path/flatbuffers.jar \ +# ./run-delta-regression.sh 'org.apache.spark.sql.delta.DeletionVectorsSuite' [...] +# +# Env: +# DELTA_VERSION delta tag to test against (default 3.3.2) +# COMET_JARS comma-separated jars added to the test classpath (required) +# JAVA_HOME JDK for sbt (17 recommended) +set -euo pipefail + +DELTA_VERSION="${DELTA_VERSION:-3.3.2}" +WORKDIR="${1:?usage: run-delta-regression.sh [...suites]}" +shift +[ $# -ge 1 ] || { echo "no suites given" >&2; exit 2; } +: "${COMET_JARS:?COMET_JARS must list the comet jars}" + +# Resolve to an absolute path before the cd below: the log path is built from +# $WORKDIR after we're already inside the Delta checkout, so a relative +# argument would otherwise be re-anchored under $DELTA_DIR. +mkdir -p "$WORKDIR" +WORKDIR=$(cd "$WORKDIR" && pwd) + +# Canonicalize each entry to an absolute path before the cd below, for the same +# reason as the WORKDIR normalization above: the injected sbt `file(p)` resolves +# a relative COMET_EXTRA_JARS entry beneath $DELTA_DIR, not the caller's directory, +# once we've already changed into the Delta checkout. +IFS=',' read -ra _jars <<< "$COMET_JARS" +_jars_abs=() +for j in "${_jars[@]}"; do + [ -f "$j" ] || { echo "COMET_JARS entry not found: $j" >&2; exit 2; } + _jars_abs+=("$(cd "$(dirname "$j")" && pwd)/$(basename "$j")") +done +COMET_JARS=$(IFS=','; echo "${_jars_abs[*]}") + +DELTA_DIR="$WORKDIR/delta-$DELTA_VERSION" +if [ ! -d "$DELTA_DIR" ]; then + git clone --depth 1 --branch "v$DELTA_VERSION" https://github.com/delta-io/delta.git "$DELTA_DIR" +elif [ ! -d "$DELTA_DIR/.git" ]; then + echo "stale/partial checkout at $DELTA_DIR; remove it (rm -rf) and rerun" >&2 + exit 2 +fi +cd "$DELTA_DIR" + +# Add COMET_EXTRA_JARS to every project's test classpath, plus the JDK-17 +# module-access flags Spark needs (both for forked test JVMs and sbt's own JVM). +if ! grep -q "COMET_EXTRA_JARS" build.sbt; then + python3 - <<'EOF' +s = open('build.sbt').read() +marker = 'lazy val commonSettings = Seq(' +opens = [ + "--add-opens=java.base/java.lang=ALL-UNNAMED", + "--add-opens=java.base/java.lang.invoke=ALL-UNNAMED", + "--add-opens=java.base/java.lang.reflect=ALL-UNNAMED", + "--add-opens=java.base/java.io=ALL-UNNAMED", + "--add-opens=java.base/java.net=ALL-UNNAMED", + "--add-opens=java.base/java.nio=ALL-UNNAMED", + "--add-opens=java.base/java.util=ALL-UNNAMED", + "--add-opens=java.base/java.util.concurrent=ALL-UNNAMED", + "--add-opens=java.base/java.util.concurrent.atomic=ALL-UNNAMED", + "--add-opens=java.base/jdk.internal.ref=ALL-UNNAMED", + "--add-opens=java.base/sun.nio.ch=ALL-UNNAMED", + "--add-opens=java.base/sun.nio.cs=ALL-UNNAMED", + "--add-opens=java.base/sun.security.action=ALL-UNNAMED", + "--add-opens=java.base/sun.util.calendar=ALL-UNNAMED", + "--add-exports=java.base/sun.nio.ch=ALL-UNNAMED", +] +opts = ", ".join('"%s"' % o for o in opens) +inject = ( + 'lazy val commonSettings = Seq(\n' + ' Test / unmanagedJars ++= sys.env.get("COMET_EXTRA_JARS").toSeq\n' + ' .flatMap(_.split(",")).map(p => Attributed.blank(file(p))),\n' + ' Test / fork := true,\n' + ' Test / javaOptions ++= Seq(%s),\n' % opts +) +assert marker in s, 'commonSettings marker not found' +open('build.sbt', 'w').write(s.replace(marker, inject, 1)) +EOF +fi + +# Inject Comet into the shared test SparkSession when COMET_EXTRA_JARS is set. +TEST_BASE=spark/src/test/scala/org/apache/spark/sql/delta/test/DeltaSQLCommandTest.scala +if ! grep -q "CometSparkSessionExtensions" "$TEST_BASE"; then + python3 - "$TEST_BASE" <<'EOF' +import sys +p = sys.argv[1] +s = open(p).read() +old = ''' override protected def sparkConf: SparkConf = { + super.sparkConf + .set(StaticSQLConf.SPARK_SESSION_EXTENSIONS.key, + classOf[DeltaSparkSessionExtension].getName) + .set(SQLConf.V2_SESSION_CATALOG_IMPLEMENTATION.key, + classOf[DeltaCatalog].getName) + }''' +new = ''' override protected def sparkConf: SparkConf = { + val conf = super.sparkConf + .set(StaticSQLConf.SPARK_SESSION_EXTENSIONS.key, + classOf[DeltaSparkSessionExtension].getName) + .set(SQLConf.V2_SESSION_CATALOG_IMPLEMENTATION.key, + classOf[DeltaCatalog].getName) + if (sys.env.contains("COMET_EXTRA_JARS")) { + conf + .set(StaticSQLConf.SPARK_SESSION_EXTENSIONS.key, + classOf[DeltaSparkSessionExtension].getName + + ",org.apache.comet.CometSparkSessionExtensions") + .set("spark.comet.enabled", "true") + .set("spark.comet.exec.enabled", "true") + .set("spark.comet.exec.shuffle.enabled", "true") + .set("spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", "2g") + .set("spark.comet.scan.delta.enabled", "true") + } else conf + }''' +assert old in s, 'sparkConf block not found' +open(p, 'w').write(s.replace(old, new)) +EOF +fi + +# ScanReportHelper is a test-only trait that counts scans by pattern-matching +# FileSourceScanExec in the executed plan. The Comet Delta scan replaces those +# nodes, so claimed scans would go uncounted ("0 did not equal 2" in +# MergeIntoSuiteBase's insert-only data-skipping test). Map the Comet node back +# to the FileSourceScanExec it was built from: originalPlan carries the same +# PreparedDeltaFileIndex, so the reported paths and skipping stats are identical. +SCAN_HELPER=spark/src/test/scala/org/apache/spark/sql/delta/test/ScanReportHelper.scala +if [ -f "$SCAN_HELPER" ] && ! grep -q "CometDeltaNativeScanExec" "$SCAN_HELPER"; then + python3 - "$SCAN_HELPER" <<'EOF' +import sys +p = sys.argv[1] +s = open(p).read() +old = " case fs: FileSourceScanExec => Seq(fs)\n" +new = (" case fs: FileSourceScanExec => Seq(fs)\n" + " case c: org.apache.spark.sql.comet.CometDeltaNativeScanExec =>\n" + " Seq(c.originalPlan)\n") +assert s.count(old) == 1, s.count(old) +open(p, 'w').write(s.replace(old, new)) +EOF +fi + +export COMET_EXTRA_JARS="$COMET_JARS" +export SPARK_LOCAL_IP=127.0.0.1 +export RUST_BACKTRACE=1 + +cmds=() +for sel in "$@"; do + cmds+=("spark/testOnly $sel") +done + +LOG="$WORKDIR/delta-regression-$(date +%Y%m%d-%H%M%S).log" +echo "==> logging to $LOG" +build/sbt "${cmds[@]}" 2>&1 | tee "$LOG" | grep -E "^\[info\] (Tests:|Suites:|All tests|.*\*\*\* FAILED| - )" | tail -80 diff --git a/contrib/delta-spark/pom.xml b/contrib/delta-spark/pom.xml new file mode 100644 index 00000000000..172a04ed27f --- /dev/null +++ b/contrib/delta-spark/pom.xml @@ -0,0 +1,241 @@ + + + + + 4.0.0 + + org.apache.datafusion + comet-parent-spark${spark.version.short}_${scala.binary.version} + 1.2.0-SNAPSHOT + ../../pom.xml + + + comet-contrib-delta-spark${spark.version.short}_${scala.binary.version} + comet-contrib-delta + + + + ${project.basedir}/../../native/target/debug + false + + + + + org.apache.datafusion + comet-spark-spark${spark.version.short}_${scala.binary.version} + ${project.version} + provided + + + io.delta + ${delta.artifact}_${scala.binary.version} + ${delta.spark.version} + provided + + + + commons-logging + commons-logging + + + + + org.apache.spark + spark-sql_${scala.binary.version} + provided + + + + com.google.flatbuffers + flatbuffers-java + 25.2.10 + test + + + + org.apache.arrow + arrow-vector + ${arrow.version} + test + + + org.apache.arrow + arrow-memory-unsafe + ${arrow.version} + test + + + org.apache.arrow + arrow-c-data + ${arrow.version} + test + + + + org.apache.parquet + parquet-column + + + org.apache.parquet + parquet-hadoop + + + + org.apache.datafusion + comet-spark-spark${spark.version.short}_${scala.binary.version} + ${project.version} + test-jar + test + + + org.scalatest + scalatest_${scala.binary.version} + test + + + + org.testcontainers + minio + + + software.amazon.awssdk + s3 + + + + org.apache.spark + spark-hadoop-cloud_${scala.binary.version} + tests + + + + com.google.guava + guava + ${guava.version} + test + + + org.scalatestplus + junit-4-13_${scala.binary.version} + test + + + org.apache.spark + spark-sql_${scala.binary.version} + ${spark.version} + test-jar + test + + + org.apache.spark + spark-core_${scala.binary.version} + ${spark.version} + test-jar + test + + + + commons-logging + commons-logging + + + + + org.apache.spark + spark-catalyst_${scala.binary.version} + ${spark.version} + test-jar + test + + + + + + + net.alchim31.maven + scala-maven-plugin + + + org.scalatest + scalatest-maven-plugin + + + org.apache.maven.plugins + maven-enforcer-plugin + ${maven-enforcer-plugin.version} + + + + no-duplicate-declared-dependencies + + + + + + org.apache.datafusion + comet-common-spark${spark.version.short}_${scala.binary.version} + + + org.apache.comet.* + + + + + + + + + + + + + + + + release + + ${project.basedir}/../../native/target/release + + + + + diff --git a/contrib/delta-spark/src/main/resources/META-INF/services/org.apache.comet.CometConfigProvider b/contrib/delta-spark/src/main/resources/META-INF/services/org.apache.comet.CometConfigProvider new file mode 100644 index 00000000000..6db01e4b245 --- /dev/null +++ b/contrib/delta-spark/src/main/resources/META-INF/services/org.apache.comet.CometConfigProvider @@ -0,0 +1,17 @@ +# 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. +org.apache.comet.contrib.delta.DeltaSparkConfigProvider diff --git a/contrib/delta-spark/src/main/resources/META-INF/services/org.apache.comet.rules.CometScanContrib b/contrib/delta-spark/src/main/resources/META-INF/services/org.apache.comet.rules.CometScanContrib new file mode 100644 index 00000000000..25a913e0cd0 --- /dev/null +++ b/contrib/delta-spark/src/main/resources/META-INF/services/org.apache.comet.rules.CometScanContrib @@ -0,0 +1,17 @@ +# 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. +org.apache.comet.contrib.delta.DeltaScanContrib diff --git a/contrib/delta-spark/src/main/resources/META-INF/services/org.apache.spark.sql.comet.PlanDataInjector b/contrib/delta-spark/src/main/resources/META-INF/services/org.apache.spark.sql.comet.PlanDataInjector new file mode 100644 index 00000000000..c26629ec377 --- /dev/null +++ b/contrib/delta-spark/src/main/resources/META-INF/services/org.apache.spark.sql.comet.PlanDataInjector @@ -0,0 +1,17 @@ +# 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. +org.apache.spark.sql.comet.DeltaPlanDataInjector diff --git a/contrib/delta-spark/src/main/scala/org/apache/comet/contrib/delta/CometDeltaNativeScan.scala b/contrib/delta-spark/src/main/scala/org/apache/comet/contrib/delta/CometDeltaNativeScan.scala new file mode 100644 index 00000000000..3659ab4c33e --- /dev/null +++ b/contrib/delta-spark/src/main/scala/org/apache/comet/contrib/delta/CometDeltaNativeScan.scala @@ -0,0 +1,599 @@ +/* + * 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.contrib.delta + +import scala.jdk.CollectionConverters._ + +import org.apache.hadoop.fs.Path +import org.apache.spark.internal.Logging +import org.apache.spark.sql.catalyst.expressions.Literal +import org.apache.spark.sql.comet.{CometScanExec, DeltaPlanDataInjector} +import org.apache.spark.sql.delta.DeltaParquetFileFormat +import org.apache.spark.sql.delta.RowIndexFilterType +import org.apache.spark.sql.delta.actions.DeletionVectorDescriptor +import org.apache.spark.sql.execution.{FileSourceScanExec, ScalarSubquery => ExecScalarSubquery} +import org.apache.spark.sql.execution.datasources.{FilePartition, PartitionedFile} +import org.apache.spark.sql.types.{ByteType, LongType, MetadataBuilder, StructField, StructType} + +import org.apache.comet.CometSparkSessionExtensions.withFallbackReason +import org.apache.comet.objectstore.NativeConfig +import org.apache.comet.serde.OperatorOuterClass +import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.QueryPlanSerde.{exprToProto, serializeDataType} +import org.apache.comet.serde.operator.{literalToProto, partition2Proto, schema2Proto, CometNativeScan} +import org.apache.comet.shims.ShimFileFormat + +/** + * Serde for the native Delta scan. Two shapes: + * - Plain reads reuse core's `NativeScanCommon` builder wholesale. + * - Deletion-vector reads: Delta's planner appends `__delta_internal_is_row_deleted` (tinyint) + * and Spark's row-index temp column (bigint) to the read schema and filters on is_row_deleted + * above the scan. The native reader applies the DV as a row selection, so both internal + * columns are emitted as per-file constants (0), the parquet read schema is stripped to the + * real data columns, and the DV descriptor ships per file for native to fetch and decode. + */ +object CometDeltaNativeScan + extends Logging + with org.apache.spark.sql.catalyst.expressions.PredicateHelper { + + val IsRowDeletedColumn: String = DeltaParquetFileFormat.IS_ROW_DELETED_COLUMN_NAME + val RowIndexColumn: String = ShimFileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME + + private[delta] val internalColumnNames: Set[String] = Set(IsRowDeletedColumn, RowIndexColumn) + + // Prefix for the internal columns' slots in the partition schema, mirroring core's + // _comet_metadata_ prefix rationale: DataFusion matches partition columns by name. + // [[allocateUniqueInternalFields]] additionally suffixes on collision with a real column. + private val deltaConstFieldPrefix = "_comet_delta_" + + def isDvShape(scanExec: FileSourceScanExec): Boolean = + scanExec.requiredSchema.exists(f => internalColumnNames.contains(f.name)) + + private def deltaFormat(scanExec: FileSourceScanExec): DeltaParquetFileFormat = + scanExec.relation.fileFormat.asInstanceOf[DeltaParquetFileFormat] + + private def columnMappingMode(scanExec: FileSourceScanExec): String = + deltaFormat(scanExec).metadata.columnMappingMode.name + + /** + * Under column mapping, parquet files store physical column names (stable UUIDs / ids), so the + * schemas passed to the native parquet reader must be physical. Positions and structure are + * preserved, so output binding and projection are unaffected. The scan's internal DV columns + * are not part of the table schema and must be stripped before calling this. + * + * `private[delta]` (not `private`): [[DeltaScanSupport.declineReason]]'s non-ASCII + * case-insensitive name gate reuses this exact conversion to compute the names native sees + * under column mapping, rather than re-deriving physical names with separate logic. + */ + private[delta] def toPhysical(scanExec: FileSourceScanExec, schema: StructType): StructType = { + val format = deltaFormat(scanExec) + if (format.metadata.columnMappingMode.name == "none") { + schema + } else { + // Name mode matches file columns by physical NAME. Strip the parquet.field.id metadata + // createPhysicalSchema also stamps: files written before the column-mapping upgrade have + // no field ids and would fail the reader's id expectations. + stripFieldIds(org.apache.spark.sql.delta.DeltaColumnMapping + .createPhysicalSchema(schema, format.metadata.schema, format.metadata.columnMappingMode)) + } + } + + private def stripFieldIds(schema: StructType): StructType = { + import org.apache.spark.sql.types._ + def stripType(dt: DataType): DataType = dt match { + case s: StructType => stripFieldIds(s) + case a: ArrayType => a.copy(elementType = stripType(a.elementType)) + case m: MapType => + m.copy(keyType = stripType(m.keyType), valueType = stripType(m.valueType)) + case other => other + } + StructType(schema.fields.map { f => + val metadata = new MetadataBuilder() + .withMetadata(f.metadata) + .remove("parquet.field.id") + // Sibling key Delta stamps on array/map fields under IcebergCompat/Uniform. + .remove("parquet.field.nested.ids") + .build() + f.copy(dataType = stripType(f.dataType), metadata = metadata) + }) + } + + /** + * Build the planning-time `DeltaScan` operator (common data only; file partitions are injected + * lazily at execution). Returns None when an output data type cannot be serialized or the plan + * shape is not one we can translate faithfully, and tags `scanExec` with the reason when the + * data files need more than one ABFS store. `memo` is the same claim-memo instance + * [[DeltaScanSupport.declineReason]] populated on this claim; its `hadoopConf` and + * `dvDescriptors` are reused here rather than recomputed. + */ + def convert( + scanExec: FileSourceScanExec, + scanHelper: CometScanExec, + memo: DeltaScanSupport.DeltaClaimMemo): Option[Operator] = { + val relation = scanExec.relation + + val dataFileUris = scanHelper.selectedPartitions.view.flatMap(_.files).map(_.getPath.toUri) + val firstFileUri = dataFileUris.headOption + + val hadoopConf = memo.hadoopConf + + val objectStoreOptions = + NativeConfig.extractScanObjectStoreOptions(hadoopConf, dataFileUris) match { + case Right(options) => options + case Left(reason) => + // declineReason already declines data files on more than one store; this only + // catches ABFS authorities that differ in case or percent-encoding alone. + withFallbackReason(scanExec, reason) + return None + } + + val tableRootPath = relation.location.rootPaths.head + val tableRoot = tableRootPath.toString + + val commonOpt = if (!isDvShape(scanExec)) { + // Under column mapping (name mode) the parquet reader must see physical names; + // positions are preserved so output binding and projection stay untouched. + CometNativeScan.buildNativeScanCommon( + source = scanExec.simpleStringWithNodeId(), + output = scanExec.output, + requiredSchema = toPhysical(scanExec, scanExec.requiredSchema), + dataSchema = toPhysical(scanExec, relation.dataSchema), + partitionSchema = toPhysical(scanExec, relation.partitionSchema), + fileConstantMetadataColumns = scanExec.fileConstantMetadataColumns, + dataFilters = scanHelper.supportedDataFilters, + objectStoreOptions = objectStoreOptions, + hadoopConf = hadoopConf, + conf = scanExec.conf) + } else { + buildDvScanCommon(scanExec, scanHelper, objectStoreOptions, hadoopConf) + } + + commonOpt.map { commonBuilder => + // Already forced by declineReason on this claim; reused rather than deserialized again. + val dvDescriptors = memo.dvDescriptors + // Union object-store options over every authority a partition of this scan may need a + // store for, not just the first data file's scheme. + commonBuilder.putAllObjectStoreOptions( + mergedObjectStoreOptions( + hadoopConf, + storeUris(dvDescriptors, tableRootPath, firstFileUri)).asJava) + + val common = commonBuilder.build() + // Effective session rebase read modes, resolved through ParquetOptions exactly as + // ParquetFileFormat.buildReaderWithPartitionValues resolves them (per-relation + // `datetimeRebaseMode` / `int96RebaseMode` options win over the session conf, whose + // per-Spark-version default -- EXCEPTION on 3.x, CORRECTED on 4.0 -- SQLConf supplies). + // Native consults them only for files whose footer metadata does not decide the rebase + // policy on its own, mirroring DataSourceUtils.getRebaseSpec's modeByConfig fallback. + val parquetReadOptions = + new org.apache.spark.sql.execution.datasources.parquet.ParquetOptions( + relation.options, + scanExec.conf) + val deltaCommon = OperatorOuterClass.DeltaSparkScanCommon + .newBuilder() + .setTableRoot(tableRoot) + .setColumnMappingMode(columnMappingMode(scanExec)) + .setSourceKey(DeltaPlanDataInjector.sourceKey(tableRoot, common)) + .setDatetimeRebaseModeInRead(parquetReadOptions.datetimeRebaseModeInRead) + .setInt96RebaseModeInRead(parquetReadOptions.int96RebaseModeInRead) + .build() + val deltaScan = OperatorOuterClass.DeltaSparkScan + .newBuilder() + .setCommon(common) + .setDeltaCommon(deltaCommon) + Operator + .newBuilder() + .setPlanId(scanExec.id) + .setContribScan(DeltaSparkScanEnvelope.pack(deltaScan.build())) + .build() + } + } + + /** + * One representative store URI per distinct object-store authority this scan's partitions may + * need options for: the data-file authority (`firstFileUri`), the table root unconditionally + * (UUID-relative DV sidecars resolve against it), and every distinct on-disk DV authority from + * `descriptors` (inline DVs carry no external URI and are filtered out). Deduping by authority + * rather than full URI keeps this O(distinct authorities) instead of O(files), keeping the + * FIRST URI seen per authority so `firstFileUri`/the table root win over a same-authority DV + * path. + */ + private[delta] def storeUris( + descriptors: Seq[DeletionVectorDescriptor], + tableRootPath: Path, + firstFileUri: Option[java.net.URI]): Seq[java.net.URI] = { + val dvAuthorityUris = descriptors + .filter(_.storageType != DeletionVectorDescriptor.INLINE_DV_MARKER) + .map(_.absolutePath(tableRootPath).toUri) + val candidates = firstFileUri.toSeq ++ Seq(tableRootPath.toUri) ++ dvAuthorityUris + val byAuthority = scala.collection.mutable.LinkedHashMap.empty[String, java.net.URI] + candidates.foreach(uri => + byAuthority.getOrElseUpdate(DeltaScanSupport.uriAuthority(uri), uri)) + byAuthority.values.toSeq + } + + /** + * Unions `NativeConfig.extractObjectStoreOptions` over every `uris` authority. Safe to union + * rather than pick one: S3 and GCS options sit under their own scheme's prefix, and + * re-extracting the same scheme from two URIs is idempotent. Azure options are not scoped to an + * account and overlap (WASB copies `fs.azure.*`; ABFS forwards `comet.azure.*` markers plus + * global `fs.azure.*` keys), but only the table root can bring an Azure authority here: + * `DeltaScanSupport.declineReason` declines Azure data files and deletion vectors, whose + * container is userinfo. + */ + private[delta] def mergedObjectStoreOptions( + hadoopConf: org.apache.hadoop.conf.Configuration, + uris: Seq[java.net.URI]): Map[String, String] = + uris.foldLeft(Map.empty[String, String]) { (merged, uri) => + merged ++ NativeConfig.extractObjectStoreOptions(hadoopConf, uri) + } + + /** + * Harvest subquery-bearing predicates for this scan from its covering FilterExec. Spark 3.x + * strips them from a scan's `dataFilters` at planning (`FileSourceStrategy` routes them to the + * post-scan filter only), while Spark 4.x keeps them in `dataFilters`; collecting them here at + * claim time gives the execution-time resolve-and-push path the same inputs on every version, + * and the dedup below keeps Spark 4.x from carrying duplicates. Reference containment alone + * does not prove pushing a predicate down is semantics-preserving, so `spineToScan` also + * requires every intervening operator to commute with the push (an intervening + * LIMIT/Sort/Aggregate/join etc. stops the walk and leaves the filter where Spark placed it: + * missed pruning only). + */ + def subqueryFiltersFromParent( + plan: org.apache.spark.sql.execution.SparkPlan, + scanExec: FileSourceScanExec): Seq[org.apache.spark.sql.catalyst.expressions.Expression] = { + import org.apache.spark.sql.catalyst.expressions.{PlanExpression, SubqueryExpression} + import org.apache.spark.sql.execution.{FilterExec, ProjectExec, SparkPlan} + + // Whether every node from `node` down to `scanExec` is one pushdown can safely cross: a + // deterministic ProjectExec is 1:1 on rows and a deterministic FilterExec only removes rows, + // so moving a predicate over the scan's output through either preserves semantics -- mirroring + // Spark's own PushPredicateThroughNonJoin/CollapseProject rules. A nondeterministic node (or + // anything else: LIMIT/TopN, Sort, Aggregate, Window, joins, ...) can change which rows survive + // to matter, so it stops the walk and the filter is left uncollected (missed pruning only). + def spineToScan(node: SparkPlan): Boolean = node match { + case n if n eq scanExec => true + case p: ProjectExec if p.projectList.forall(_.deterministic) => spineToScan(p.child) + case f: FilterExec if f.condition.deterministic => spineToScan(f.child) + case _ => false + } + + // Nearest FilterExec whose spine down to the scan is Project/Filter-only (the DV shape + // interposes such nodes between them, so do not require a direct parent-child edge). + val filtersAboveScan = plan.collect { + case f: FilterExec if spineToScan(f.child) => f + } + filtersAboveScan.lastOption + .map { f => + splitConjunctivePredicates(f.condition) + .filter(_.deterministic) + .filter(_.references.subsetOf(scanExec.outputSet)) + .filter(p => + SubqueryExpression.hasSubquery(p) || p.exists(_.isInstanceOf[PlanExpression[_]])) + .filterNot(p => scanExec.dataFilters.exists(_.semanticEquals(p))) + } + .getOrElse(Seq.empty) + } + + /** + * Execution-time scalar-subquery data filters of a scan. `hasResolvedFilters` is true whenever + * pushdown is enabled and such filters exist, whether or not they bind or serialize; `protos` + * holds only the ones that serialized. + */ + case class ResolvedSubqueryFilters( + hasResolvedFilters: Boolean, + protos: Seq[org.apache.comet.serde.ExprOuterClass.Expr]) + + private val NoResolvedSubqueryFilters = ResolvedSubqueryFilters(false, Seq.empty) + + /** + * Resolve scalar-subquery data filters at execution time and serialize them for native + * pushdown, mirroring `CometNativeScanExec.serializedPartitionData`. `supportedDataFilters` + * excludes PlanExpressions at planning time (subquery results do not exist yet), so these + * bounds reach the native reader only through this path. Filters that fail to serialize are + * skipped: Spark keeps a covering FilterExec above the scan, so this is missed pruning only. + * Their presence is still reported, since native keys the safe timestamp conversion on the scan + * being filtered at all, as the core scan does for its resolved filters. + * + * Known core-parity limitation: when fused under a parent native operator, + * `ensureSubqueriesResolved` has already called `updateResult()` on these subqueries and this + * path calls it again (`ScalarSubquery.updateResult` re-executes unconditionally); benign here + * since the subquery's snapshot is pinned at analysis, but wasteful. Fix belongs in core. + */ + def resolvedSubqueryFilters( + dataFilters: Seq[org.apache.spark.sql.catalyst.expressions.Expression], + output: Seq[org.apache.spark.sql.catalyst.expressions.Attribute], + requiredSchema: StructType, + conf: org.apache.spark.sql.internal.SQLConf): ResolvedSubqueryFilters = { + if (!conf.getConf(org.apache.spark.sql.internal.SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED)) { + return NoResolvedSubqueryFilters + } + val subqueryFilters = dataFilters.filter(_.exists(_.isInstanceOf[ExecScalarSubquery])) + if (subqueryFilters.isEmpty) { + return NoResolvedSubqueryFilters + } + // Same binding guard as the DV shape's plan-time filters: references limited to the + // data-column prefix of the output, where positions agree with the native read schema. + // Guard BEFORE updateResult so discarded filters never execute their subqueries. + val strippedLen = requiredSchema.count(f => !internalColumnNames.contains(f.name)) + val dataColIds = output.take(strippedLen).map(_.exprId).toSet + val pushableFilters = + subqueryFilters.filter(_.references.forall(r => dataColIds.contains(r.exprId))) + pushableFilters.foreach(_.foreach { + case s: ExecScalarSubquery => s.updateResult() + case _ => + }) + val protos = pushableFilters + .flatMap { filter => + // MergeScalarSubqueries can fuse several scalar subqueries into one struct-returning + // subquery accessed via GetStructField; fold that whole subtree to a literal (a bare + // GetStructField-over-Literal would not serialize). + val resolved = filter.transform { + case g @ org.apache.spark.sql.catalyst.expressions + .GetStructField(_: ExecScalarSubquery, _, _) => + Literal.create(g.eval(null), g.dataType) + case s: ExecScalarSubquery => + Literal.create(s.eval(null), s.dataType) + } + val proto = exprToProto(resolved, output) + if (proto.isEmpty) { + logWarning(s"Could not serialize resolved scalar subquery filter: $resolved") + } + proto + } + ResolvedSubqueryFilters(hasResolvedFilters = true, protos) + } + + /** + * Allocate the partition-schema slots for the DV shape's internal columns + * (`internalColumnNames`), with names collision-free against the physical data schema, the + * physical partition schema, and the constant-metadata slots already allocated for this scan + * (plus each other): DataFusion substitutes partition constants BY NAME, so an unprefixed, + * un-uniquified slot could collide with a real column and silently replace its data with the + * bookkeeping constant. `buildDvScanCommon` keys `internalIndexByName` by each field's ORIGINAL + * name from `requiredSchema`, so the renaming here only changes the proto's field name. + */ + private[delta] def allocateUniqueInternalFields( + requiredSchema: StructType, + physicalDataSchema: StructType, + physicalPartitionSchema: StructType, + constantMetadataFields: Seq[StructField]): Seq[StructField] = { + val reserved = scala.collection.mutable.LinkedHashSet[String]() + reserved ++= physicalDataSchema.fields.map(_.name) + reserved ++= physicalPartitionSchema.fields.map(_.name) + reserved ++= constantMetadataFields.map(_.name) + requiredSchema.fields.toSeq + .filter(f => internalColumnNames.contains(f.name)) + .map { f => + var name = s"$deltaConstFieldPrefix${f.name}" + while (reserved.contains(name)) { + name = name + "_" + } + reserved += name + StructField(name, f.dataType, f.nullable) + } + } + + /** + * DV shape common builder. Layout invariants (declined by DeltaScanSupport when violated): scan + * output = requiredSchema attrs (data columns, then the internal columns as a suffix) followed + * by partition and constant-metadata columns. The parquet read schema strips the internal + * columns; they are appended to the partition schema as per-file constants, so the projection + * vector routes them from the constants block. + */ + private def buildDvScanCommon( + scanExec: FileSourceScanExec, + scanHelper: CometScanExec, + objectStoreOptions: Map[String, String], + hadoopConf: org.apache.hadoop.conf.Configuration) + : Option[OperatorOuterClass.NativeScanCommon.Builder] = { + val relation = scanExec.relation + val output = scanExec.output + val requiredSchema = scanExec.requiredSchema + + val commonBuilder = OperatorOuterClass.NativeScanCommon.newBuilder() + commonBuilder.setSource(scanExec.simpleStringWithNodeId()) + + val scanTypes = output.flatMap(attr => serializeDataType(attr.dataType)) + if (scanTypes.length != output.length) { + return None + } + commonBuilder.addAllFields(scanTypes.asJava) + + val strippedRequired = + StructType(requiredSchema.filterNot(f => internalColumnNames.contains(f.name))) + val strippedLen = strippedRequired.length + val requiredLen = requiredSchema.length + + // Keep only data filters that bind identically in the output and the native index space: + // references limited to the first strippedLen output attributes. Internal-column filters + // (is_row_deleted = 0) are trivially true after native DV application. + if (scanExec.conf.getConf( + org.apache.spark.sql.internal.SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED)) { + commonBuilder.setHasDataFilters(scanHelper.supportedDataFilters.nonEmpty) + val dataColIds = output.take(strippedLen).map(_.exprId).toSet + val filterProtos = scanHelper.supportedDataFilters + .filter(_.references.forall(r => dataColIds.contains(r.exprId))) + .flatMap(f => exprToProto(f, output)) + commonBuilder.addAllDataFilters(filterProtos.asJava) + } + + // Real partition columns carry physical names in the proto, same as the data/required + // schemas: a retained physical data name can otherwise collide with a partition column's + // LOGICAL name after a rename history, and DataFusion's by-name partition rewrite would then + // replace the data projection with the partition constant. constantMetadataFields/ + // internalFields are synthetic slots, not table columns, so they are not physicalized. + val physicalDataSchema = toPhysical(scanExec, relation.dataSchema) + val physicalPartitionSchema = toPhysical(scanExec, relation.partitionSchema) + // Constant metadata and real partition columns follow the required schema in the output, + // exactly like the plain shape. Names are uniquified against the physical data/partition + // schemas for the same by-name-collision reason [[allocateUniqueInternalFields]] exists. + val constantMetadataFields = CometNativeScan.uniqueConstantMetadataFields( + scanExec.fileConstantMetadataColumns, + physicalDataSchema.fields.map(_.name).toSet ++ physicalPartitionSchema.fields + .map(_.name) + .toSet) + val internalFields = allocateUniqueInternalFields( + requiredSchema, + physicalDataSchema = physicalDataSchema, + physicalPartitionSchema = physicalPartitionSchema, + constantMetadataFields = constantMetadataFields) + val partitionSchemaFields = + physicalPartitionSchema.fields.toSeq ++ constantMetadataFields ++ internalFields + + // Protos carry physical names (column mapping); index math below stays logical. + val partitionSchemaProto = schema2Proto(partitionSchemaFields) + val physicalRequired = toPhysical(scanExec, strippedRequired) + val requiredSchemaProto = schema2Proto(physicalRequired) + val dataSchemaProto = schema2Proto(physicalDataSchema) + + // Projection: data columns from the (stripped) read schema; internal columns from their + // constants slots at the END of the partition fields; the output tail (real partitions + + // constant metadata) positionally from the head of the partition fields. + val dataSchema = relation.dataSchema + val internalBase = dataSchema.length + partitionSchemaFields.length - internalFields.length + val internalIndexByName = requiredSchema.fields.toSeq + .filter(f => internalColumnNames.contains(f.name)) + .zipWithIndex + .map { case (f, i) => f.name -> (internalBase + i) } + .toMap + val projectionVector = output.zipWithIndex.map { case (attr, i) => + val idx = if (internalColumnNames.contains(attr.name)) { + internalIndexByName(attr.name) + } else if (i < requiredLen) { + dataSchema.fieldIndex(attr.name) + } else { + dataSchema.length + (i - requiredLen) + } + idx.toLong.asInstanceOf[java.lang.Long] + } + commonBuilder.addAllProjectionVector(projectionVector.asJava) + + commonBuilder.addAllDataSchema(dataSchemaProto.asJava) + commonBuilder.addAllRequiredSchema(requiredSchemaProto.asJava) + commonBuilder.addAllPartitionSchema(partitionSchemaProto.asJava) + + // The physical schema, as in the plain shape: it is what DeltaParquetFileFormat hands + // Spark's ParquetReadSupport, so the field id flags match the ids Spark checks. + CometNativeScan.populateScanConfFlags( + commonBuilder, + physicalRequired, + objectStoreOptions, + hadoopConf, + scanExec.conf) + + Some(commonBuilder) + } + + /** Serialize one file partition into a DeltaSparkScan proto with per-file DV descriptors. */ + def serializePartition( + filePartition: FilePartition, + scanExec: FileSourceScanExec, + tableRoot: String): Array[Byte] = { + val relation = scanExec.relation + val sparkPartition = partition2Proto( + filePartition, + relation.partitionSchema, + scanExec.fileConstantMetadataColumns, + ShimFileFormat.fileConstantMetadataExtractors(relation.fileFormat)) + + val dvShape = isDvShape(scanExec) + + val deltaPartition = OperatorOuterClass.DeltaSparkFilePartition.newBuilder() + sparkPartition.getPartitionedFileList.asScala.zip(filePartition.files.toSeq).foreach { + case (fileProto, file) => + val fileBuilder = fileProto.toBuilder + if (dvShape) { + // Append the internal-constant values after the real partition/constant-metadata + // values, matching the order of the appended partition-schema fields. + scanExec.requiredSchema.fields + .filter(f => internalColumnNames.contains(f.name)) + .foreach { f => + val lit = f.dataType match { + case ByteType => Literal(0.toByte, ByteType) + case LongType => Literal(0L, LongType) + case other => + // Fixed internal invariant (observed Delta 3.3 types); fail loudly on + // drift rather than emit a plausible-looking constant. + throw new IllegalStateException( + s"Unexpected type $other for Delta internal column ${f.name}") + } + fileBuilder.addPartitionValues( + literalToProto(lit, s"delta internal constant ${f.name}")) + } + } + val dfb = OperatorOuterClass.DeltaSparkPartitionedFile + .newBuilder() + .setFile(fileBuilder.build()) + extractDvDescriptor(file, tableRoot).foreach(dfb.setDv) + deltaPartition.addPartitionedFile(dfb.build()) + } + + OperatorOuterClass.DeltaSparkScan + .newBuilder() + .setFilePartition(deltaPartition.build()) + .build() + .toByteArray + } + + /** + * Pull the DV descriptor Delta attached to this file (base64 under + * `row_index_filter_id_encoded`), resolving UUID-relative paths to absolute URLs and + * Z85-decoding inline bitmaps here on the JVM where delta-spark's codecs live. + */ + private def extractDvDescriptor( + file: PartitionedFile, + tableRoot: String): Option[OperatorOuterClass.DeltaSparkDvDescriptor] = { + val encoded = file.otherConstantMetadataColumnValues + .get(DeltaParquetFileFormat.FILE_ROW_INDEX_FILTER_ID_ENCODED) + val filterType = file.otherConstantMetadataColumnValues + .get(DeltaParquetFileFormat.FILE_ROW_INDEX_FILTER_TYPE) + encoded.map { enc => + filterType match { + case Some(RowIndexFilterType.IF_CONTAINED) | None => + case other => + // DeltaScanSupport declines CDF reads, the only source of inverted filters; + // reaching here means a gate was bypassed -- fail loudly rather than corrupt. + throw new IllegalStateException( + s"Native Delta scan cannot apply row index filter type $other") + } + val desc = DeletionVectorDescriptor.deserializeFromBase64(enc.asInstanceOf[String]) + val builder = OperatorOuterClass.DeltaSparkDvDescriptor + .newBuilder() + .setStorageType(desc.storageType) + .setSizeInBytes(desc.sizeInBytes) + .setCardinality(desc.cardinality) + if (desc.storageType == DeletionVectorDescriptor.INLINE_DV_MARKER) { + // Delegates to core, which owns the shaded/relocated dependency this field's setter + // is generated against, so this module's source never has to name that package. + CometNativeScan.setDvInlineData(builder, desc.inlineData) + } else { + // Same convention as data-file paths (SparkPath.urlEncoded): a raw Hadoop path + // with spaces or % characters would be mangled by the native URL parse. + builder.setAbsolutePath( + org.apache.spark.paths.SparkPath + .fromPath(desc.absolutePath(new Path(tableRoot))) + .urlEncoded) + desc.offset.foreach(builder.setOffset) + } + builder.build() + } + } +} diff --git a/contrib/delta-spark/src/main/scala/org/apache/comet/contrib/delta/DeltaScanConf.scala b/contrib/delta-spark/src/main/scala/org/apache/comet/contrib/delta/DeltaScanConf.scala new file mode 100644 index 00000000000..c94e01493dc --- /dev/null +++ b/contrib/delta-spark/src/main/scala/org/apache/comet/contrib/delta/DeltaScanConf.scala @@ -0,0 +1,76 @@ +/* + * 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.contrib.delta + +import org.apache.comet.{ConfigBuilder, ConfigEntry} + +/** + * Configuration for the JVM-planned Delta Lake scan contrib. The support is experimental and + * explicitly opt-in: having the contrib jar on the classpath is not enough, the scan must also be + * enabled with `spark.comet.scan.delta.enabled`. + * + * This is the plain, user-facing flag for enabling native Delta scans, kept under the + * `spark.comet.scan.delta` namespace. The experimental Rust-kernel-backed scan path is a separate + * opt-in, defined by the kernel contrib's own `DeltaConf` under the + * `spark.comet.scan.deltaNative` namespace; the two jars define distinct keys and are not + * expected to coexist -- see the ownership contract in `CometScanContrib`. Entry construction + * self-registers with `CometConf.allConfs` via the `ConfigBuilder` machinery. + */ +object DeltaScanConf { + + // Matches the kernel contrib's category so both group onto the same generated-docs table. + private[delta] val CATEGORY = "delta" + + val COMET_DELTA_NATIVE_ENABLED: ConfigEntry[Boolean] = + ConfigBuilder("spark.comet.scan.delta.enabled") + .category(CATEGORY) + .doc( + "Whether to enable native Delta table scans. When enabled, DSv1 Delta table reads " + + "planned by delta-spark are executed through Comet's native Parquet scan, " + + "inheriting row-group pruning, page-index pruning, and filter pushdown, with " + + "deletion vectors applied inside the scan. Experimental: defaults to false, so " + + "adding the contrib jar does not by itself change how any query is read.") + .booleanConf + .createWithDefault(false) + + val COMET_DELTA_MAX_DELETED_ROWS_PER_FILE: ConfigEntry[Long] = + ConfigBuilder("spark.comet.scan.delta.dv.maxDeletedRowsPerFile") + .category(CATEGORY) + .doc( + "Upper bound on a single file's deletion-vector cardinality (deleted row count) the " + + "native Delta scan will claim. Applying a deletion vector expands it into per-row " + + "selectors that are held in memory. This bound caps one file's selectors, not what a " + + "task holds: the selectors for every file in a partition stay held until the task " + + "finishes. The bound is a deliberately pessimistic planning-time proxy for that " + + "memory (deletion vector cardinality, not the exact selector count), so a large but " + + "contiguous deletion is declined the same as a large alternating one. Scans whose " + + "deletion vectors exceed this bound for any file fall back to Spark's reader.") + .longConf + .createWithDefault(1000000) + + /** + * Every entry defined here, in docs order. Referencing this forces object initialisation, which + * registers the entries -- see `CometConfigProvider`. + */ + def all: Seq[ConfigEntry[_]] = + Seq(COMET_DELTA_MAX_DELETED_ROWS_PER_FILE, COMET_DELTA_NATIVE_ENABLED) + + def scanEnabled: Boolean = COMET_DELTA_NATIVE_ENABLED.get() +} diff --git a/contrib/delta-spark/src/main/scala/org/apache/comet/contrib/delta/DeltaScanContrib.scala b/contrib/delta-spark/src/main/scala/org/apache/comet/contrib/delta/DeltaScanContrib.scala new file mode 100644 index 00000000000..f34bcaec110 --- /dev/null +++ b/contrib/delta-spark/src/main/scala/org/apache/comet/contrib/delta/DeltaScanContrib.scala @@ -0,0 +1,107 @@ +/* + * 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.contrib.delta + +import org.apache.spark.internal.Logging +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.comet.CometDeltaNativeScanExec +import org.apache.spark.sql.execution.{FileSourceScanExec, SparkPlan} +import org.apache.spark.sql.execution.datasources.HadoopFsRelation + +import org.apache.comet.CometConf.COMET_EXEC_ENABLED +import org.apache.comet.CometSparkSessionExtensions.{hasFallbackReason, withFallbackReason} +import org.apache.comet.rules.CometScanContrib + +/** + * Claims DSv1 Delta Lake scans for native execution, discovered by core's ServiceLoader (see + * `META-INF/services/org.apache.comet.rules.CometScanContrib`). Scans this contrib owns but + * cannot handle are claimed with a tagged fallback reason (per the `CometScanContrib` ownership + * contract); Spark's Delta reader then handles them. + * + * The produced `CometDeltaNativeScanExec` is fully converted at claim time, so this contrib does + * NOT use `CometContribScanMarker` (which exists for planning-time nodes that `CometExecRule` + * converts later; mixing it in here would convert the node a second time). + */ +class DeltaScanContrib extends CometScanContrib with Logging { + + override def tryTransformV1( + plan: SparkPlan, + session: SparkSession, + scanExec: FileSourceScanExec, + relation: HadoopFsRelation): Option[SparkPlan] = { + // Not a Delta scan: not ours; core handles it exactly as before. + if (!DeltaScanSupport.isDeltaScan(scanExec)) { + return None + } + + // Contrib scans are native-exec nodes, so like core's own nativeScan they require + // COMET_EXEC_ENABLED. Our old core-side hook gated all extensions centrally; the + // CometScanContrib call site does not, so the gate lives here. Silent None (no tag) + // preserves the old "never consulted" behavior and avoids double-tagging next to + // core's own exec-disabled fallback reason. + if (!COMET_EXEC_ENABLED.get()) { + return None + } + + if (!DeltaScanConf.scanEnabled) { + // Deliberate deviation from the "own but cannot handle => claim" contract: a + // user-disabled contrib must be fully inert (the jar alone changes nothing) and must + // not shadow another registered Delta contrib. Tag the opt-in hint for EXPLAIN, pass. + withFallbackReason( + scanExec, + "Native Delta scan not enabled: set " + + s"${DeltaScanConf.COMET_DELTA_NATIVE_ENABLED.key}=true to opt in") + return None + } + + // Built before declineReason (rather than only on claim) so the multi-object-store gate + // can inspect the scan's selected files without listing them twice; convert reuses this + // same helper on a claim. + val scanHelper = + CometDeltaNativeScanExec.planningHelper(scanExec, scanExec.partitionFilters) + // Populated by declineReason on the claimable path only, and reused by convert below so a + // claimed scan does not recompute the Hadoop conf or the DV descriptors a second time. + val claimMemo = new DeltaScanSupport.DeltaClaimMemo + DeltaScanSupport.declineReason(plan, scanExec, scanHelper, claimMemo) match { + case Some(reason) => + Some(withFallbackReason(scanExec, reason)) + case None => + CometDeltaNativeScan.convert(scanExec, scanHelper, claimMemo) match { + case Some(nativeOp) => + logDebug( + s"COMET-DELTA-CLAIM required=${scanExec.requiredSchema.map(_.name).mkString(",")} " + + s"output=${scanExec.output.map(_.name).mkString(",")} " + + s"dvShape=${CometDeltaNativeScan.isDvShape(scanExec)} " + + s"planRoot=${plan.getClass.getSimpleName}") + val subqueryDataFilters = + CometDeltaNativeScan.subqueryFiltersFromParent(plan, scanExec) + Some(CometDeltaNativeScanExec(nativeOp, scanExec, subqueryDataFilters)) + // convert tags its own reason when it declines the scan's object stores. + case None if hasFallbackReason(scanExec) => + Some(scanExec) + case None => + Some( + withFallbackReason( + scanExec, + "Native Delta scan does not support the scan's output data types")) + } + } + } +} diff --git a/contrib/delta-spark/src/main/scala/org/apache/comet/contrib/delta/DeltaScanSupport.scala b/contrib/delta-spark/src/main/scala/org/apache/comet/contrib/delta/DeltaScanSupport.scala new file mode 100644 index 00000000000..91b9006af6e --- /dev/null +++ b/contrib/delta-spark/src/main/scala/org/apache/comet/contrib/delta/DeltaScanSupport.scala @@ -0,0 +1,2026 @@ +/* + * 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.contrib.delta + +import java.io.IOException +import java.net.{URI, URISyntaxException} +import java.util.Locale + +import scala.collection.mutable.{ListBuffer, Map => MutableMap} +import scala.jdk.CollectionConverters._ + +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.Path +import org.apache.spark.sql.catalyst.expressions.{Alias, ExprId, GenericInternalRow, InputFileBlockLength, InputFileBlockStart, InputFileName} +import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, GenericArrayData} +import org.apache.spark.sql.catalyst.util.ResolveDefaultColumns.getExistenceDefaultValues +import org.apache.spark.sql.comet.CometScanExec +import org.apache.spark.sql.delta.DeltaParquetFileFormat +import org.apache.spark.sql.delta.actions.DeletionVectorDescriptor +import org.apache.spark.sql.execution.{FileSourceScanExec, ProjectExec, SparkPlan} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{ArrayType, DataType, MapType, StructType} + +import org.apache.comet.CometConf +import org.apache.comet.CometConf.COMET_LIBHDFS_SCHEMES +import org.apache.comet.objectstore.NativeConfig +import org.apache.comet.parquet.CometParquetUtils +import org.apache.comet.rules.{CometScanRule, CometScanTypeChecker} +import org.apache.comet.serde.operator.CometNativeScan +import org.apache.comet.shims.ShimFileFormat + +/** + * Claim/decline gates for the native Delta scan. Correctness rule: when in doubt, decline, + * Spark's Delta reader handles the scan and results stay correct, just unaccelerated. + */ +object DeltaScanSupport { + + /** + * Reader features the native path understands; anything else on the protocol declines the + * table. `deletionVectors`/`columnMapping` are declined separately below for specific reasons. + */ + private val understoodReaderFeatures: Set[String] = + Set("columnMapping", "deletionVectors", "timestampNtz", "v2Checkpoint", "vacuumProtocolCheck") + + /** + * Is this exactly Delta's DSv1 parquet format? Compared by class name, not `classOf`: a + * `classOf` reference would raise `NoClassDefFoundError` and break every parquet scan when + * delta-spark is absent from the classpath. + */ + def isDeltaScan(scanExec: FileSourceScanExec): Boolean = + scanExec.relation.fileFormat.getClass.getName == + "org.apache.spark.sql.delta.DeltaParquetFileFormat" + + /** + * Claim-time artifacts [[declineReason]] already computes but [[CometDeltaNativeScan.convert]] + * also needs -- threaded through by reference (populated only on the claimable path, right + * before `declineReason` returns `None`) so a claimed scan does not pay to recompute either: + * the Hadoop conf ([[org.apache.spark.sql.internal.SessionState#newHadoopConfWithOptions]] is + * not cheap) and the deletion-vector descriptors (base64-decoded, non-trivial only for DV-shape + * scans). One instance is created per claim attempt in `DeltaScanContrib` and passed to both + * `declineReason` and `convert`. + */ + private[delta] final class DeltaClaimMemo { + var hadoopConf: Configuration = _ + var dvDescriptors: Seq[DeletionVectorDescriptor] = Seq.empty + } + + /** + * First reason this Delta scan cannot go native, or None when claimable (in which case `memo` + * is populated for [[CometDeltaNativeScan.convert]] to reuse). Only called when [[isDeltaScan]] + * is true. `scanHelper` is the [[CometScanExec]] built to drive `convert` on a claim, reused + * for the multi-store gate below. + */ + def declineReason( + plan: SparkPlan, + scanExec: FileSourceScanExec, + scanHelper: CometScanExec, + memo: DeltaClaimMemo): Option[String] = { + val format = scanExec.relation.fileFormat.asInstanceOf[DeltaParquetFileFormat] + val protocol = format.protocol + val metadata = format.metadata + // Name mode is supported via physical-name schemas; id mode needs the field-id path and + // stays declined until validated. Hoisted here since several gates below reuse it. + val cmMode = metadata.columnMappingMode.name + // Descriptor deserialization is expensive, so hoist it into a `lazy val`, forced at most + // once in this method; on the claimable path the result is handed to `convert` through + // `memo` below, so a claimed scan deserializes the descriptors exactly once end to end. + val tableRoot = scanExec.relation.location.rootPaths.head.toString + lazy val dvDescriptors: Seq[DeletionVectorDescriptor] = + selectedDvDescriptors(scanHelper, tableRoot) + + // Mirrors core's CometScanRule.isSchemaSupported so scan-time type gates (unsigned-small-int + // fallback, collation, shredded-variant-struct) apply identically here. Pure in-memory check, + // so it runs first, ahead of every I/O-bearing gate below. + // Unlike core, a required Variant root stays declined: this path lacks core's Variant gates. + val schemaFallbackReasons = new ListBuffer[String]() + val typeChecker = CometScanTypeChecker() + val requiredSchemaSupported = + typeChecker.isSchemaSupported(scanExec.requiredSchema, schemaFallbackReasons) + val partitionSchemaSupported = + typeChecker.isSchemaSupported(scanExec.relation.partitionSchema, schemaFallbackReasons) + if (!requiredSchemaSupported || !partitionSchemaSupported) { + return Some( + "Native Delta scan does not support the schema: " + schemaFallbackReasons.mkString(", ")) + } + + if (format.isCDCRead) { + return Some("Native Delta scan does not support Change Data Feed reads") + } + + // Delta's DML machinery (findTouchedFiles) disables reader optimizations and needs real + // row indexes from Spark's reader; claiming here would feed NULL indexes into DV construction. + if (!format.optimizationsEnabled) { + return Some("Native Delta scan does not support reads with reader optimizations disabled") + } + if (scanExec.requiredSchema.exists(_.name == DeltaParquetFileFormat.ROW_INDEX_COLUMN_NAME) || + scanExec.relation.dataSchema.exists( + _.name == DeltaParquetFileFormat.ROW_INDEX_COLUMN_NAME)) { + return Some("Native Delta scan does not support Delta's generated row-index column") + } + + if (cmMode != "none" && cmMode != "name") { + return Some(s"Native Delta scan does not support column mapping mode $cmMode") + } + // createPhysicalSchema wholesale-replaces field metadata, silently dropping EXISTS_DEFAULT. + if (cmMode == "name" && + getExistenceDefaultValues(scanExec.requiredSchema).exists(_ != null)) { + return Some( + "Native Delta scan does not support column defaults together with column mapping") + } + // createPhysicalSchema rewrites nested StructField names too, and the native builder emits the + // required schema verbatim as output, so name-sensitive expressions (e.g. to_json) would leak + // physical names. Decline until a rename adapter exists. + if (cmMode == "name" && + scanExec.requiredSchema.exists(f => containsNestedStruct(f.dataType))) { + return Some("Native Delta scan does not support column mapping with nested struct fields") + } + + val readerFeatures = protocol.readerFeatureNames + val unknownFeatures = readerFeatures -- understoodReaderFeatures + if (unknownFeatures.nonEmpty) { + return Some( + s"Native Delta scan does not support reader feature(s) ${unknownFeatures.mkString(", ")}") + } + + // file_block_start and file_block_length describe the split reading a row. For a split file, + // DataFusion keeps a row group in the split holding its first page while Spark's reader + // keeps it in the split holding its midpoint, so rows would report the wrong split. Core + // declines the same two columns (https://github.com/apache/datafusion-comet/issues/6505). + // A deletion-vector read carries Delta's whole _metadata struct, so the two columns can sit + // in the scan output unused; only a value consumed above the scan declines. + val consumedSplitDependentCols = scanExec.fileConstantMetadataColumns + .filter(a => Set("file_block_start", "file_block_length").contains(a.name)) + .filterNot(a => attributesUnusedAbove(plan, scanExec, Set(a.exprId))) + if (consumedSplitDependentCols.nonEmpty) { + return Some( + "Native Delta scan does not support metadata column(s) " + + consumedSplitDependentCols.map(_.name).mkString(", ")) + } + + // Non-constant metadata columns are generated per-row by Spark's reader and unsupported, + // except Delta's DV bookkeeping columns, which the native path emits as constants. + val knownColNames = + scanExec.relation.dataSchema.map(_.name).toSet ++ + scanExec.relation.partitionSchema.map(_.name).toSet ++ + scanExec.fileConstantMetadataColumns.map(_.name).toSet ++ + CometDeltaNativeScan.internalColumnNames + val unknownOutput = scanExec.output.map(_.name).filterNot(knownColNames.contains) + if (unknownOutput.nonEmpty) { + return Some( + s"Native Delta scan does not support generated column(s) ${unknownOutput.mkString(", ")}") + } + + // Deletion-vector shape invariants (see CometDeltaNativeScan.buildDvScanCommon). + if (CometDeltaNativeScan.isDvShape(scanExec)) { + // A row-index column WITHOUT is_row_deleted is Delta DML bookkeeping (real row indexes), + // not a DV read; claiming it with a constant would corrupt the DVs being written. + val hasIsRowDeleted = + scanExec.requiredSchema.exists(_.name == CometDeltaNativeScan.IsRowDeletedColumn) + val hasRowIndex = + scanExec.requiredSchema.exists(_.name == CometDeltaNativeScan.RowIndexColumn) + if (hasRowIndex && !hasIsRowDeleted) { + return Some( + "Native Delta scan does not support row-index reads outside a deletion-vector scan") + } + // Internal columns must form a suffix of the read schema so data-column positions agree + // between Spark's output and the stripped native schema. + val names = scanExec.requiredSchema.fields.map(_.name) + val firstInternal = names.indexWhere(CometDeltaNativeScan.internalColumnNames.contains) + if (!names.drop(firstInternal).forall(CometDeltaNativeScan.internalColumnNames.contains)) { + return Some("Native Delta scan requires DV bookkeeping columns to trail the read schema") + } + // Native applies the DV itself and emits a dead constant for row-index, so the real value + // must be provably unused above the scan. + if (!rowIndexUnusedAbove(plan, scanExec)) { + return Some( + "Native Delta scan cannot supply _metadata.row_index values consumed by the query") + } + // The DV common builder does not serialize existence defaults yet. + if (getExistenceDefaultValues(scanExec.requiredSchema).exists(_ != null)) { + return Some( + "Native Delta scan does not support column defaults together with deletion vectors") + } + // Bounds native's memory for expanded DV row selectors (delta_dv.rs), pessimistically + // bounded by 2*cardinality + #row-groups; the conf below makes an over-pessimistic decline + // recoverable. + val maxDeletedRowsPerFile = DeltaScanConf.COMET_DELTA_MAX_DELETED_ROWS_PER_FILE.get() + val oversizedCardinalities = dvDescriptors + .map(_.cardinality) + .filter(_ > maxDeletedRowsPerFile) + if (oversizedCardinalities.nonEmpty) { + return Some( + "Native Delta scan does not support a deletion vector deleting " + + s"${oversizedCardinalities.max} rows in a single file, exceeding " + + s"${DeltaScanConf.COMET_DELTA_MAX_DELETED_ROWS_PER_FILE.key}=$maxDeletedRowsPerFile") + } + } + + // input_file_name & friends read from a thread-local Spark's FileScanRDD sets; the native scan + // does not populate it, and Delta's DML find-touched-files scans use it (mirrors core's check + // in CometScanRule.nativeScan). + if (plan.exists(node => + node.expressions.exists(_.exists { + case _: InputFileName | _: InputFileBlockStart | _: InputFileBlockLength => true + case _ => false + }))) { + return Some( + "Native Delta scan is not compatible with input_file_name, " + + "input_file_block_start, or input_file_block_length") + } + + // Row-index metadata columns are generated per-row by Spark's reader (mirrors core); the DV + // shape's trailing row-index column is exempt since the gates above already proved it dead. + if (!CometDeltaNativeScan.isDvShape(scanExec) && + ShimFileFormat.findRowIndexColumnIndexInSchema(scanExec.requiredSchema) >= 0) { + return Some("Native Delta scan does not support row index generation") + } + + // Mirror core's vectorized-reader compatibility gate. + if (!SQLConf.get.getConf(SQLConf.PARQUET_VECTORIZED_READER_ENABLED) && + !CometConf.COMET_SCAN_ALLOW_DISABLED_PARQUET_VECTORIZED_READER.get()) { + return Some( + "Native Delta scan is incompatible with " + + s"${SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key}=false") + } + + // Decline ALL encrypted-parquet configurations (stricter than core): the exec node does not + // yet wire the decryption-key broadcast to executors. + val hadoopConf = scanExec.relation.sparkSession.sessionState + .newHadoopConfWithOptions(scanExec.relation.options) + // Populated now (rather than only at the very end) so it is available even though several + // early-return gates below still lie ahead: cheap to set, and every one of those gates + // declines the scan anyway, so `memo` is simply never read by `convert` in that case. + memo.hadoopConf = hadoopConf + if (CometParquetUtils.encryptionEnabled(hadoopConf)) { + return Some("Native Delta scan does not support encrypted parquet") + } + + // Nested-type column defaults cannot be serialized; a dropped default would misalign the + // value/index lists consumed positionally on the native side. Mirrors core's + // transformV1Scan gate. + val possibleDefaultValues = getExistenceDefaultValues(scanExec.requiredSchema) + if (possibleDefaultValues.exists(d => + d != null && (d.isInstanceOf[ArrayBasedMapData] || d + .isInstanceOf[GenericInternalRow] || d.isInstanceOf[GenericArrayData]))) { + return Some("Native Delta scan does not support default values for nested types") + } + + // An opted-in S3-compliant alias scheme (fs.comet.s3Compliant.schemes) is declined before the + // generic scheme gate below so the reason says why: core's native scan reads it through the + // S3 client, but the S3 divergence gates further down model Hadoop's S3AFileSystem only. + val rootUris = scanExec.relation.location.rootPaths.map(_.toUri) + val aliasReason = s3CompliantAliasSchemeReason(hadoopConf, rootUris) + if (aliasReason.isDefined) { + return aliasReason + } + + // Only claim scans whose root paths object_store (or the configured libhdfs schemes) can + // actually read (mirrors core's unsupportedFsSchemes gate). + val libhdfs = libhdfsSchemes + val unsupportedRootSchemes = unsupportedSchemes(rootUris, libhdfs) + if (unsupportedRootSchemes.nonEmpty) { + return Some( + "Native Delta scan does not support filesystem scheme(s) " + + s"${unsupportedRootSchemes.mkString(", ")}") + } + + // A recognized scheme can still carry a path object_store rejects (a directory name with a + // newline surfaces as `%0A`), which native planning hard-fails on while Spark's reader opens + // it. Mirrors core's root-path gate; the complete selected paths are probed below. + val rejectedRoot = objectStoreRejectedPathReason(rootUris, libhdfs) + if (rejectedRoot.isDefined) { + return rejectedRoot + } + + // A shallow clone can span multiple object-store authorities, but the native builder resolves + // ObjectStoreUrl from only the FIRST selected file; force file listing and decline rather than + // risk reading a later file through the wrong handle. + val dataFileUris = + scanHelper.selectedPartitions.iterator.flatMap(_.files).map(_.getPath.toUri).toSeq + + // Both gates below need the DV absolute-path URIs; dvDescriptors is already memoized. + val dvUris = dvDescriptors + .filter(_.storageType != DeletionVectorDescriptor.INLINE_DV_MARKER) + .map(_.absolutePath(new Path(tableRoot)).toUri) + + // The root-path gate above only inspects the table root(s); selected files can resolve + // through a different scheme (e.g. `viewfs:`). Checked before the authority gates below, + // which presume every URI is natively resolvable. + val unsupportedSelected = unsupportedSelectedSchemeReason(dataFileUris ++ dvUris, libhdfs) + if (unsupportedSelected.isDefined) { + return unsupportedSelected + } + + // Same path probe for every complete selected path, not just its directory: a shallow + // clone's source can sit outside this root, and CONVERT TO DELTA keeps the source Parquet + // basenames, so the rejected character can be in the file name itself. The probe is a + // native URL parse with no I/O, so once per distinct URI costs less than the scan's own + // per-file parse. + val rejectedSelected = objectStoreRejectedPathReason(dataFileUris ++ dvUris, libhdfs) + if (rejectedSelected.isDefined) { + return rejectedSelected + } + + // Checked before multiStoreReason, which presumes every URI resolves to a single store + // identity -- a userinfo-bearing authority provably does not (store keying drops userinfo). + val userInfoReason = userInfoBearingAuthorityReason(dataFileUris ++ dvUris) + if (userInfoReason.isDefined) { + return userInfoReason + } + + val multiStore = multiStoreReason(dataFileUris) + if (multiStore.isDefined) { + return multiStore + } + + // GCS's zero-I/O, conf-only credential-forwarding gate; ordered alongside the S3 credential + // gates below since all presume a single, well-formed store identity per URI. + val gcsAuthReason = gcsHadoopOnlyAuthReason(hadoopConf, dataFileUris ++ dvUris) + if (gcsAuthReason.isDefined) { + return gcsAuthReason + } + + // Zero-I/O, conf-only, like the GCS gate above: decline any bucket configured for an + // encryption algorithm outside the allowlist (SSE-C, CSE-KMS, CSE-CUSTOM, or unknown) before + // the credential-divergence gates below, which do not otherwise notice this table is readable + // through Hadoop only because Hadoop's request factory (SSE-C) or SDK-level decryption layer + // (CSE-*) does something native never learns about. + // + // Shared across this gate and the ones below: propagateBucketOptions is a full Configuration + // deep copy, and each gate would otherwise recompute it independently for the same bucket(s). + // One cache, populated lazily per bucket on first use, makes it a single copy per bucket. + val propagatedConfCache = MutableMap.empty[String, Configuration] + val encryptionReason = unsupportedEncryptionAlgorithmReason( + hadoopConf, + dataFileUris ++ dvUris, + propagatedConfCache) + if (encryptionReason.isDefined) { + return encryptionReason + } + + // Always zero-I/O (plain propagated-conf read, no keystore): native's S3 client has no + // HTTP proxy support at all (no fs.s3a.proxy.* key is read anywhere in s3.rs), so a bucket + // requiring a proxy for S3 egress must decline here rather than claim and then connect + // directly, bypassing whatever network-segmentation/firewall policy required the proxy. + val proxyReason = proxyGateReason(hadoopConf, dataFileUris ++ dvUris, propagatedConfCache) + if (proxyReason.isDefined) { + return proxyReason + } + + // Zero-I/O, conf-only, like the proxy gate above: Hadoop's AssumedRoleCredentialProvider + // sends fs.s3a.assumed.role.policy as the session policy of its STS AssumeRole request, + // while native's assumed-role provider never reads the key -- a claimed scan would assume + // the role WITHOUT the configured session restriction, silently widening permissions. + val rolePolicyReason = + assumedRolePolicyGateReason(hadoopConf, dataFileUris ++ dvUris, propagatedConfCache) + if (rolePolicyReason.isDefined) { + return rolePolicyReason + } + + // Zero-I/O, conf-only: Hadoop prefixes a scheme-less fs.s3a.endpoint with http:// when + // SSL is disabled, while native always prefixes https://, so the two sides would talk to + // different endpoints. Same shape for the STS endpoint keys, which native never reads. + val endpointReason = + hadoopOnlyEndpointGateReason(hadoopConf, dataFileUris ++ dvUris, propagatedConfCache) + if (endpointReason.isDefined) { + return endpointReason + } + + // Every fs.s3a.* option native's get_config (s3.rs) resolves must agree between what Hadoop + // itself would use and what native would read from the forwarded, substituted conf (covers + // long-form bucket credentials, JCEKS/credential-provider shadowing, and any other + // short-vs-effective divergence in one mechanism); reuses hadoopConf from the encryption gate + // above. + val s3Reason = + s3ConfigDivergenceReason(hadoopConf, dataFileUris ++ dvUris, propagatedConfCache) + if (s3Reason.isDefined) { + return s3Reason + } + + // A credential-provider class native's build_aws_credential_provider_metadata (s3.rs) does + // not recognize errors at scan EXECUTION time, after the scan was already claimed; decline + // eagerly instead. + val providerReason = providerClassGateReason(hadoopConf, dataFileUris ++ dvUris) + if (providerReason.isDefined) { + return providerReason + } + + // Reuse core's generic native-scan gates (ignoreCorruptFiles/ignoreMissingFiles, AQE DPP on + // Spark 3.4, exec enabled, existence default values, the Variant read confs, and a proto + // representation for every serialized data and partition type); tags its own fallback + // reasons. + if (!CometNativeScan.isSupported(scanExec)) { + return Some("Core native scan gates rejected the scan (see reasons above)") + } + + // Claimable: hand the already-forced descriptors to `convert` via `memo` so it does not + // deserialize them a second time. + memo.dvDescriptors = dvDescriptors + None + } + + /** + * Deletion-vector descriptors for every file this DV-shape scan selected, normalized to + * absolute on-disk paths. Returns `Seq.empty` for the plain shape. Shared by the DV cardinality + * gate and [[CometDeltaNativeScan.convert]]'s object-store option merge. + */ + private[delta] def selectedDvDescriptors( + scanHelper: CometScanExec, + tableRoot: String): Seq[DeletionVectorDescriptor] = { + if (!CometDeltaNativeScan.isDvShape(scanHelper.wrapped)) { + return Seq.empty + } + val tableRootPath = new Path(tableRoot) + scanHelper.selectedPartitions.iterator + .flatMap(_.files) + .flatMap { file => + file.metadata + .get(DeltaParquetFileFormat.FILE_ROW_INDEX_FILTER_ID_ENCODED) + .map(enc => DeletionVectorDescriptor.deserializeFromBase64(enc.asInstanceOf[String])) + } + .map(_.copyWithAbsolutePath(tableRootPath)) + .toSeq + } + + /** + * The libhdfs scheme exemption set from [[org.apache.comet.CometConf.COMET_LIBHDFS_SCHEMES]], + * parsed exactly like core's scan gate (`NativeConfig.parseSchemeSet`: split on commas, + * trimmed, lowercased) and defaulting to `Set("hdfs")` when unset. + */ + private[delta] def libhdfsSchemes: Set[String] = COMET_LIBHDFS_SCHEMES.get() match { + case Some(s) => NativeConfig.parseSchemeSet(s) + case None => Set("hdfs") + } + + /** + * Decline reason when any of `uris` uses a scheme opted in as an S3-compliant alias through + * `fs.comet.s3Compliant.schemes` (e.g. `blob`), or `None`. Core's native Parquet scan admits + * such a scheme and reads it through its S3 client, with `NativeConfig` translating the vendor + * `fs...*` keys into `fs.s3a.bucket.*` options. Spark, however, reads the + * same table through the vendor's own Hadoop FileSystem, not `S3AFileSystem`, and every S3 + * divergence gate in this object ([[s3ConfigDivergenceReason]] and its siblings) is verified + * against `S3AFileSystem`'s consumers only. With no model of how the vendor filesystem resolves + * its configuration, whether native and Spark would agree cannot be decided, so the scan is + * declined rather than claimed on a guess. Selected data-file and deletion-vector URIs under an + * alias scheme are declined by the generic scheme gates, which never admit an alias (see + * [[unsupportedSchemes]]). + */ + private[delta] def s3CompliantAliasSchemeReason( + hadoopConf: Configuration, + uris: Seq[URI]): Option[String] = { + val aliases = NativeConfig.resolveS3CompliantSchemes(hadoopConf) + if (aliases.isEmpty) { + return None + } + val found = uris + .flatMap(uri => Option(uri.getScheme).map(_.toLowerCase(Locale.ROOT))) + .filter(aliases.contains) + .distinct + if (found.isEmpty) { + None + } else { + Some( + "Native Delta scan does not support S3-compliant alias filesystem scheme(s) " + + s"${found.sorted.mkString(", ")} (${CometConf.COMET_S3_COMPLIANT_SCHEMES_KEY}): " + + "Spark reads them through a vendor filesystem whose S3 configuration resolution the " + + "native scan's S3AFileSystem divergence model cannot verify") + } + } + + /** + * The lowercased, deduplicated schemes among `uris` that neither `libhdfs` nor Comet's native + * object_store layer ([[CometScanRule.isNativelyReadableScheme]]) can read. A `null` scheme is + * tolerated, not flagged, since such a URI cannot come from a Hadoop-backed source. The alias + * set handed to core's gate is deliberately empty: an `fs.comet.s3Compliant.schemes` alias is + * never admitted here (see [[s3CompliantAliasSchemeReason]]), even though core admits it. + */ + private[delta] def unsupportedSchemes(uris: Seq[URI], libhdfs: Set[String]): Set[String] = { + uris + .filter { uri => + val sch = uri.getScheme + sch != null && { + val sl = sch.toLowerCase(Locale.ROOT) + !libhdfs.contains(sl) && !CometScanRule.isNativelyReadableScheme(uri, Set.empty) + } + } + .map(_.getScheme.toLowerCase(Locale.ROOT)) + .toSet + } + + /** + * Decline reason naming the first of `uris` whose path object_store rejects even though it + * recognizes the scheme ([[CometScanRule.objectStoreAcceptsPath]], e.g. a directory name + * containing a newline, `%0A` in the URI), or `None`. Schemes in `libhdfs` never reach + * object_store's path parser and are skipped, as is a `null` scheme (see + * [[unsupportedSchemes]]); an S3-compliant alias is declined before this gate runs. The probe + * is uncached but is a plain native URL parse with no I/O + * ([[CometScanRule.objectStoreAcceptsPath]]), so callers pass every complete selected path + * (root paths, data files and deletion vectors) once per distinct URI; a converted table can + * carry the rejected character in a file basename. The reason masks any userinfo in the named + * URI ([[redactedAuthority]]). + */ + private[delta] def objectStoreRejectedPathReason( + uris: Seq[URI], + libhdfs: Set[String]): Option[String] = { + uris.distinct + .find { uri => + val sch = uri.getScheme + sch != null && !libhdfs.contains(sch.toLowerCase(Locale.ROOT)) && + !CometScanRule.objectStoreAcceptsPath(uri) + } + .map { uri => + // Mask userinfo (see redactedAuthority); the raw path keeps its percent encoding so the + // reason shows the rejected sequence as written. + val shown = + if (uriUserInfo(uri).isEmpty) uri.toString + else s"${redactedAuthority(uri)}${Option(uri.getRawPath).getOrElse("")}" + s"Native Delta scan cannot open path '$shown': object_store rejects it " + + "(e.g. an unsupported character in the path)" + } + } + + /** + * Decline reason when any of `uris` -- the scan's selected data-file and deletion-vector URIs + * -- use a scheme [[unsupportedSchemes]] flags, or `None` when every URI is natively readable + * (or libhdfs-exempt). + */ + private[delta] def unsupportedSelectedSchemeReason( + uris: Seq[URI], + libhdfs: Set[String]): Option[String] = { + val schemes = unsupportedSchemes(uris, libhdfs) + if (schemes.isEmpty) { + None + } else { + Some( + "Native Delta scan does not support selected data file or deletion vector filesystem " + + s"scheme(s) ${schemes.mkString(", ")}") + } + } + + /** + * Decline reason when `uris` span more than one object-store authority (scheme + lowercased raw + * authority, so e.g. `S3A://Bucket` and `s3a://bucket` collapse), or `None` when they share + * one. `file://` paths carry no authority, so local scans across many directories are + * unaffected. + */ + private[delta] def multiStoreReason(uris: Seq[URI]): Option[String] = { + val authorities = uris.map(uriAuthority).distinct + if (authorities.size > 1) { + Some( + "Native Delta scan does not support data files spanning multiple object stores " + + s"(found: ${authorities.sorted.mkString(", ")})") + } else { + None + } + } + + /** + * Normalizes `uri` to a lowercased `scheme://authority` string, keyed on the raw `getAuthority` + * rather than the parsed host/port/userinfo fields: `getHost` (and `getUserInfo`/`getPort`) + * return `null` for the whole authority when it fails RFC 3986 `reg-name` syntax (e.g. an + * underscore in a GCS bucket name, `gs://my_bucket`), which would silently collapse distinct + * buckets into one empty-host key. A `null` authority normalizes to the empty string. + */ + private[delta] def uriAuthority(uri: URI): String = { + val scheme = Option(uri.getScheme).map(_.toLowerCase(Locale.ROOT)).getOrElse("") + val authority = Option(uri.getAuthority).map(_.toLowerCase(Locale.ROOT)).getOrElse("") + s"$scheme://$authority" + } + + /** + * The raw userinfo component of `uri`'s authority, or empty when none. Splits at the LAST `@` + * rather than using `URI#getUserInfo`, which (like [[uriAuthority]]'s getters) returns `null` + * for the whole authority on an RFC 3986 `reg-name` violation. Never lowercased: userinfo is + * case-sensitive. + */ + private[delta] def uriUserInfo(uri: URI): String = { + val authority = Option(uri.getAuthority).getOrElse("") + val at = authority.lastIndexOf('@') + if (at >= 0) authority.substring(0, at) else "" + } + + /** + * Redacts `uri`'s authority to `scheme`, then `://`, then a literal `***` masking userinfo, + * then `@host[:port]`, for embedding in a decline reason. NEVER interpolate `uri.getAuthority` + * or [[uriUserInfo]] directly into a reason string: doing so would leak credentials embedded as + * URI userinfo into the SQL plan's explain output, fallback-reason logging, or the Spark UI. + */ + private[delta] def redactedAuthority(uri: URI): String = { + val scheme = Option(uri.getScheme).map(_.toLowerCase(Locale.ROOT)).getOrElse("") + val authority = Option(uri.getAuthority).getOrElse("") + val at = authority.lastIndexOf('@') + val hostPort = if (at >= 0) authority.substring(at + 1) else authority + s"$scheme://***@$hostPort" + } + + /** + * Decline reason when any of `uris` carries userinfo in its authority (e.g. the container in an + * abfss:// path), or `None` when none do. The native store cache, `ObjectStoreUrl`, and + * DataFusion registry all key on scheme/host/port only, dropping userinfo, so two authorities + * differing only in userinfo collide onto the same store handle. + */ + private[delta] def userInfoBearingAuthorityReason(uris: Seq[URI]): Option[String] = { + val offending = uris.filter(uri => uriUserInfo(uri).nonEmpty).map(redactedAuthority).distinct + if (offending.isEmpty) { + None + } else { + Some("Native Delta scan does not support object-store paths whose authority carries " + + "userinfo (e.g. the container in an abfss:// path): the native object-store cache, " + + "ObjectStoreUrl and DataFusion registry all key on scheme, host and port only, so two " + + "containers on one storage account share a single store handle " + + s"(found: ${offending.sorted.mkString(", ")})") + } + } + + /** + * String-literal Hadoop conf keys consulted below. `hadoop-aws` is NOT on this module's runtime + * classpath, so `org.apache.hadoop.fs.s3a.Constants` must never be referenced here (would raise + * `NoClassDefFoundError` for sessions with no S3 dependency). + */ + private val HadoopCredentialProviderPathKey = "hadoop.security.credential.provider.path" + private val S3aCredentialProviderPathKey = "fs.s3a.security.credential.provider.path" + + /** + * `CommonConfigurationKeysPublic.HADOOP_SECURITY_CREDENTIAL_CLEAR_TEXT_FALLBACK`, default + * `true`, verified via `javap` against `hadoop-common` 3.3.4's + * `Configuration#getPasswordFromConfig`: `getPassword` only falls back to reading a plaintext + * conf value once `getBoolean(, true)` holds -- with the flag off, a plaintext value + * is invisible to every `getPassword`-based resolver, even when no credential provider is + * configured at all. + */ + private val ClearTextFallbackKey = "hadoop.security.credential.clear-text-fallback" + + private def s3aBucketProviderPathKey(bucket: String): String = + s"fs.s3a.bucket.$bucket.security.credential.provider.path" + + /** + * The LONG form of [[s3aBucketProviderPathKey]]: `S3AUtils#lookupPassword` resolves per-bucket + * overrides through both a long key (`fs.s3a.bucket.B.`) and a short key; both + * must be covered here too. + */ + private def s3aBucketLongProviderPathKey(bucket: String): String = + s"fs.s3a.bucket.$bucket.fs.s3a.security.credential.provider.path" + + private def nonEmptyConf(hadoopConf: Configuration, key: String): Boolean = + Option(hadoopConf.get(key)).exists(_.nonEmpty) + + /** + * The lowercase-scheme-checked S3/S3A bucket name from `uri`'s authority, or `None` when + * `uri`'s scheme is not `s3`/`s3a`. Parses the raw authority manually rather than + * `URI#getHost`, avoiding the same RFC 3986 `reg-name` pitfall as [[uriAuthority]]. + */ + private def s3Bucket(uri: URI): Option[String] = { + val scheme = Option(uri.getScheme).map(_.toLowerCase(Locale.ROOT)) + if (scheme.contains("s3") || scheme.contains("s3a")) { + val authority = Option(uri.getAuthority).getOrElse("") + val at = authority.lastIndexOf('@') + val hostAndPort = if (at >= 0) authority.substring(at + 1) else authority + val colon = hostAndPort.lastIndexOf(':') + val host = if (colon >= 0) hostAndPort.substring(0, colon) else hostAndPort + if (host.isEmpty) None else Some(host) + } else { + None + } + } + + private def plainValue(hadoopConf: Configuration, key: String): Option[String] = + Option(hadoopConf.get(key)).filter(_.nonEmpty) + + /** + * How Hadoop's OWN consumer reads one of the keys compared by [[s3ConfigDivergenceReason]], + * which decides how [[s3KeyDivergenceReason]] computes the Hadoop-effective side of its + * equality check. Exactly two consumer families exist among [[AllS3ConfigKeys]] in `hadoop-aws` + * 3.3.4, each verified via `javap`/CFR against the real call sites (cited per key on + * [[S3ConfigKeyConsumers]]). The tier must mirror the key's ACTUAL consumer: resolving a + * [[PropagatedOptionConsumer]] key through the wider `lookupPassword` cascade is NOT fail-safe + * for a value-EQUALITY comparator -- a long-form alias value Hadoop itself never reads can + * EQUAL native's resolution while Hadoop's true propagate-then-plain-get value differs, turning + * a real divergence into a wrongly-claimed scan (the endpoint `${...}`-redirect shape pinned in + * `DeltaScanContribSuite`). + */ + private[delta] sealed trait S3ConfigConsumer + + /** + * Read via `S3AUtils#lookupPassword(bucket, conf, baseKey)`, verified via `javap` against + * `hadoop-aws` 3.3.4: builds `longBucketKey = "fs.s3a.bucket." + bucket + "." + baseKey` (the + * FULL, already-`fs.s3a`-prefixed base key appended after the bucket segment) and reads it via + * `Configuration#getPassword` BEFORE the short-bucket key, keeping the long value whenever + * `getPassword` returns non-empty and only falling through to short-then-global otherwise. + * `getPassword` is Hadoop-credential-provider-aware and skips plaintext conf entirely when + * [[ClearTextFallbackKey]] is false. Modeled by [[hadoopLookupPasswordEffective]]. + */ + private[delta] case object LookupPasswordConsumer extends S3ConfigConsumer + + /** + * Read via `S3AUtils#propagateBucketOptions` followed by a plain `Configuration#get`-family + * call (`getTrimmed`/`getBoolean`/`getClasses`) against the propagated view: the short bucket + * form wins only by having overwritten the global key during propagation, the long bucket form + * folds into an unread `fs.s3a.fs.s3a.*` key, and neither a credential provider nor + * [[ClearTextFallbackKey]] is ever consulted. Modeled as a plain `Configuration#get` on the + * [[propagateBucketOptions]] result, which also expands `${...}` references under that + * propagated view exactly like the real consumer. + */ + private[delta] case object PropagatedOptionConsumer extends S3ConfigConsumer + + /** + * Every `fs.s3a.*` base key that governs whether a claimed native scan actually behaves like + * Hadoop's own reader would, paired with the consumer family Hadoop resolves it through -- ONE + * list, with each key's resolution tier declared beside it, so a key can never sit in the + * comparator without a deliberate classification (adding one without picking a tier does not + * compile). The entries are every per-bucket `fs.s3a.*` base key native's S3 client's + * `get_config` (s3.rs) resolves, verified directly against its call sites: + * `extract_s3_config_options` (endpoint.region, path.style.access, endpoint, + * requester.pays.enabled), `lookup_provider_class` (the Comet-specific + * credential-provider-class activation key), and + * `build_credential_provider`/`build_aws_credential_provider_metadata`/ + * `build_assume_role_credential_provider_metadata` (aws.credentials.provider, + * assumed.role.credentials.provider, assumed.role.arn, assumed.role.session.name). + * + * Tier assignments, each verified via `javap`/CFR against `hadoop-aws` 3.3.4: + * - access.key/secret.key/session.token: `S3AUtils#getAWSAccessKeys` and + * `MarshalledCredentialBinding#fromFileSystem` (reached from + * `TemporaryAWSCredentialsProvider`) resolve all three via `S3AUtils#lookupPassword` -- + * [[LookupPasswordConsumer]]. + * - aws.credentials.provider and assumed.role.credentials.provider: + * `S3AUtils#buildAWSProviderList` -> `loadAWSProviderClasses` -> plain + * `Configuration#getClasses` -- [[PropagatedOptionConsumer]]. + * - assumed.role.arn/session.name: `AssumedRoleCredentialProvider`'s constructor reads both + * via plain `Configuration#getTrimmed` -- [[PropagatedOptionConsumer]]. + * - endpoint (`S3AFileSystem`: `getTrimmed`), endpoint.region (`DefaultS3ClientFactory`: + * `getTrimmed`), path.style.access (`S3AFileSystem`: `getBoolean`) -- + * [[PropagatedOptionConsumer]]. + * - requester.pays.enabled: not read anywhere in `hadoop-aws` 3.3.4 (the constant does not + * even exist in its `Constants` class); later releases read it via plain `getBoolean` + * against the propagated conf, so the plain tier is both the faithful forward model and + * inert on 3.3.4 -- [[PropagatedOptionConsumer]]. + * - comet.credential.provider.class: Comet's own activation key, plain conf read on both + * sides, never a Hadoop key at all -- [[PropagatedOptionConsumer]]. + * + * SYNC NOTE: the key list must stay a superset of native's `NATIVE_S3A_CONFIG_PROPERTIES` + * constant (`native/core/src/parquet/objectstore/s3.rs`, property suffixes without the + * `fs.s3a.` prefix) -- `DeltaScanContribSuite`'s discovery-harness test asserts this + * mechanically against [[AllS3ConfigKeys]]. A Rust test in `s3.rs` also parses this `Seq(` + * block as text, so each entry must stay a single `"fs.s3a." -> ` literal pair. + * Literal strings, not the [[AwsCredentialsProviderKey]] / + * [[AssumedRoleCredentialsProviderKey]] vals declared below, purely to avoid a forward + * reference inside this `object` body; kept textually identical to those two constants. + */ + private[delta] val S3ConfigKeyConsumers: Seq[(String, S3ConfigConsumer)] = Seq( + "fs.s3a.access.key" -> LookupPasswordConsumer, + "fs.s3a.secret.key" -> LookupPasswordConsumer, + "fs.s3a.session.token" -> LookupPasswordConsumer, + "fs.s3a.aws.credentials.provider" -> PropagatedOptionConsumer, + "fs.s3a.assumed.role.arn" -> PropagatedOptionConsumer, + "fs.s3a.assumed.role.session.name" -> PropagatedOptionConsumer, + "fs.s3a.assumed.role.credentials.provider" -> PropagatedOptionConsumer, + "fs.s3a.endpoint" -> PropagatedOptionConsumer, + "fs.s3a.endpoint.region" -> PropagatedOptionConsumer, + "fs.s3a.path.style.access" -> PropagatedOptionConsumer, + "fs.s3a.requester.pays.enabled" -> PropagatedOptionConsumer, + "fs.s3a.comet.credential.provider.class" -> PropagatedOptionConsumer) + + /** The compared keys alone, in [[S3ConfigKeyConsumers]] order (discovery-harness surface). */ + private[delta] val AllS3ConfigKeys: Seq[String] = S3ConfigKeyConsumers.map(_._1) + + /** + * The short-bucket-then-global value resolved for `baseKey` under `bucket` from `hadoopConf`, + * skipping an empty value at either alias exactly like [[plainValue]]. NOT used by + * [[s3ConfigDivergenceReason]]/[[s3KeyDivergenceReason]] any more -- every key checked there + * resolves per its declared [[S3ConfigKeyConsumers]] tier (see [[s3KeyDivergenceReason]]), + * neither of which matches this function's read. This function's one remaining caller is + * [[shortThenGlobalOrReason]], which reads provider-CLASS strings (from the ORIGINAL, + * unpropagated conf) for name-support validation in [[providerClassReason]]/ + * [[assumedRoleProviderClassReason]] -- by the time those run, [[s3ConfigDivergenceReason]] has + * already proven Hadoop's and native's effective values agree for the same key, so whichever of + * the two (equal) values this narrower read returns does not affect correctness there. NEVER + * used to compute native's own effective value -- see [[nativeShortThenGlobal]] for that. + */ + private def shortThenGlobal( + hadoopConf: Configuration, + bucket: String, + baseKey: String): Option[String] = { + val shortKey = s"fs.s3a.bucket.$bucket." + baseKey.stripPrefix("fs.s3a.") + plainValue(hadoopConf, shortKey).orElse(plainValue(hadoopConf, baseKey)) + } + + /** + * The short-bucket-then-global value native's `get_config` (s3.rs) resolves for `baseKey` under + * `bucket` from the ORIGINAL, unpropagated `hadoopConf` -- `NativeConfig + * .extractObjectStoreOptions` forwards `Configuration#get`'s substituted value for every + * `fs.s3a.*` entry with no bucket-option propagation step of its own, so the original conf is + * the right input here. Unlike [[shortThenGlobal]]/[[plainValue]], this mirrors `get_config` + * faithfully: PRESENCE of the short-bucket key alone -- never its emptiness -- decides whether + * native falls back to the global key (`get_config` is a plain `HashMap::get`, which returns + * `Some` for a key explicitly set to `""`), so an explicitly empty or whitespace-only + * short-bucket value resolves to `Some("")` here and never falls through to global -- the + * OPPOSITE of Hadoop's own `getPassword`/`lookupPassword` semantics (see + * [[hadoopLookupPasswordEffective]]), which treat empty as absent and keep trying the next + * alias. The ONLY function used to compute native's effective value in + * [[s3KeyDivergenceReason]]. + * + * Deliberately does NOT apply `get_config_trimmed`'s `.trim()` here: [[s3KeyDivergenceReason]] + * trims both this value and Hadoop's effective value together, symmetrically, at the point they + * are compared, rather than one-sidedly here -- trimming only the native side would flag a + * spurious divergence for a value neither side's whitespace actually changes the behavior of + * once each side's own downstream parsing normalizes it (e.g. Hadoop's own multi-line + * `fs.s3a.aws.credentials.provider` default, which both Hadoop and native additionally trim per + * comma-separated entry after splitting), while a one-sided trim would make an + * otherwise-identical default value look diverged for every bucket, never claiming natively at + * all. + */ + private def nativeShortThenGlobal( + hadoopConf: Configuration, + bucket: String, + baseKey: String): Option[String] = { + val shortKey = s"fs.s3a.bucket.$bucket." + baseKey.stripPrefix("fs.s3a.") + Option(hadoopConf.get(shortKey)).orElse(Option(hadoopConf.get(baseKey))) + } + + /** + * Faithful in-memory replica of `S3AUtils#propagateBucketOptions` (`hadoop-aws`), which + * `S3AFileSystem#initialize` calls FIRST, before any option or credential is read: + * `Configuration conf = propagateBucketOptions(originalConf, bucket); ...; setConf(conf);` -- + * every subsequent `conf.get`/`getPassword` call in that filesystem instance, including + * `${...}` variable substitution, resolves against this propagated view, not the original conf. + * `hadoop-aws` is not on this module's runtime classpath (see the string-literal-keys note + * above), so `S3AUtils#propagateBucketOptions` cannot be called directly; this reproduces its + * logic verbatim using only `hadoop-common`'s `Configuration`: + * {{{ + * public static Configuration propagateBucketOptions(Configuration source, String bucket) { + * final String bucketPrefix = FS_S3A_BUCKET_PREFIX + bucket + '.'; + * final Configuration dest = new Configuration(source); + * for (Map.Entry entry : source) { + * final String key = entry.getKey(); + * final String value = entry.getValue(); // the (unexpanded) value + * if (!key.startsWith(bucketPrefix) || bucketPrefix.equals(key)) continue; + * final String stripped = key.substring(bucketPrefix.length()); + * if (stripped.startsWith("bucket.") || "impl".equals(stripped)) { + * // ignored + * } else { + * final String generic = FS_S3A_PREFIX + stripped; + * dest.set(generic, value, ...); // overwrites any existing global value + * } + * } + * return dest; + * } + * }}} + * Note the LONG bucket form (`fs.s3a.bucket.B.fs.s3a.`) folds to an unread + * `fs.s3a.fs.s3a.` key here too, exactly like the real method -- `stripped` already starts + * with `fs.s3a.` in that case, so prepending `fs.s3a.` again produces a key nothing ever reads. + */ + private def propagateBucketOptions(hadoopConf: Configuration, bucket: String): Configuration = { + val bucketPrefix = s"fs.s3a.bucket.$bucket." + val dest = new Configuration(hadoopConf) + hadoopConf.iterator().asScala.foreach { entry => + val key = entry.getKey + if (key.startsWith(bucketPrefix) && key != bucketPrefix) { + val stripped = key.substring(bucketPrefix.length) + if (!stripped.startsWith("bucket.") && stripped != "impl") { + dest.set(s"fs.s3a.$stripped", entry.getValue) + } + } + } + dest + } + + /** + * Canonical and deprecated Hadoop S3A encryption-algorithm config keys, verified via `javap` + * against `hadoop-aws` 3.3.4's `org.apache.hadoop.fs.s3a.Constants`: `S3_ENCRYPTION_ALGORITHM = + * "fs.s3a.encryption.algorithm"` (canonical) and `SERVER_SIDE_ENCRYPTION_ALGORITHM = + * "fs.s3a.server-side-encryption-algorithm"` (DEPRECATED -- note the hyphen before "algorithm", + * unlike the corresponding `*.key` constants below, which both use a `.key` suffix). + * `hadoop-aws` is NOT on this module's runtime classpath, so these stay string literals, same + * rationale as [[HadoopCredentialProviderPathKey]]. + */ + private val S3EncryptionAlgorithmKey = "fs.s3a.encryption.algorithm" + private val DeprecatedS3EncryptionAlgorithmKey = "fs.s3a.server-side-encryption-algorithm" + + /** + * The exact strings `S3AEncryptionMethods#getMethod` accepts, verified via `javap`/CFR against + * `hadoop-aws` 3.3.4's `S3AEncryptionMethods` enum: `NONE("")`, `SSE_S3("AES256", serverSide = + * true, requiresSecret = false)`, `SSE_KMS("SSE-KMS", serverSide = true, requiresSecret = + * false)`, `SSE_C("SSE-C", serverSide = true, requiresSecret = true)`, `CSE_KMS("CSE-KMS", + * serverSide = false, requiresSecret = true)`, `CSE_CUSTOM("CSE-CUSTOM", serverSide = false, + * requiresSecret = true)`. `getMethod` parses case-insensitively + * (`values().find(_.getMethod.equalsIgnoreCase(algorithm))`), matched below the same way. + * + * ALLOWLIST, not a blocklist (replaces the former SSE-C-only blocklist): only the algorithms S3 + * decrypts transparently on GET/HEAD given read permission alone, with NO extra request header + * and NO client-side step, are safe for a native scan that forwards none of Hadoop's + * `fs.s3a.encryption.*`/`fs.s3a.server-side-encryption*` options -- + * - `AES256` (SSE_S3, `serverSide = true`): plain server-side encryption, transparent on GET. + * - `SSE-KMS` (SSE_KMS, `serverSide = true`): server-side, KMS-managed key, transparent on + * GET given KMS decrypt permission (no header). + * - `DSSE-KMS`: NOT present in this enum on `hadoop-aws` 3.3.4 (confirmed by the six values + * listed above) -- `S3AEncryptionMethods.getMethod("DSSE-KMS")` throws + * `IOException("Unknown encryption algorithm DSSE-KMS")` on this version, so + * `S3AUtils#buildEncryptionSecrets` (and therefore Hadoop's own reader) already fails + * before ever reading such a table under 3.3.4, meaning this string can never actually be + * the resolved value on the declared target version -- admitting it here is inert there. + * Included anyway, forward-compatible, for a newer `hadoop-aws` on the runtime classpath (a + * later Hadoop release; this module has no compile-time `hadoop-aws` dependency, see the + * string-literal-keys note above) where DSSE-KMS is a real, dual-layer, server-side + * algorithm decrypted transparently on GET the same way SSE-KMS is. Every other value + * declines: `SSE-C` (SSE_C is `serverSide = true` in Hadoop's own enum, but `requiresSecret + * \= true` -- S3 rejects a GET/HEAD for an SSE-C object outright (400 Bad Request) unless + * the customer key is resent as a request header on every call, so a native scan that never + * learns the key cannot succeed at all, where Hadoop's own reader -- whose request factory + * attaches the key -- would), `CSE-KMS`/`CSE-CUSTOM` (`serverSide = false`: client-side + * encryption decrypts object bytes locally in the SDK layer, which the native Parquet + * reader has no equivalent of -- it would read raw ciphertext), and any future/unknown + * value (a value `S3AEncryptionMethods.getMethod` itself would reject is certainly not one + * of the three confirmed-transparent algorithms above; declining is the only safe default + * for anything this gate cannot positively confirm). + */ + private val AllowedEncryptionAlgorithms: Set[String] = Set("AES256", "SSE-KMS", "DSSE-KMS") + + /** + * `bucket`'s effective encryption-algorithm key and value under `hadoopConf`, or `None` when + * neither the canonical nor deprecated key is set anywhere consulted. Mirrors + * `S3AUtils#buildEncryptionSecrets`'s real resolution order, verified via `javap`/CFR + * decompilation of `hadoop-aws` 3.3.4's `S3AUtils.class`: + * {{{ + * String algorithm = lookupBucketSecret(bucket, conf, "fs.s3a.encryption.algorithm"); + * if (algorithm == null) + * algorithm = lookupBucketSecret(bucket, conf, "fs.s3a.server-side-encryption-algorithm"); + * if (algorithm == null) + * algorithm = lookupPassword(null, conf, "fs.s3a.encryption.algorithm"); + * if (algorithm == null) + * algorithm = lookupPassword(null, conf, "fs.s3a.server-side-encryption-algorithm"); + * }}} + * i.e. bucket-tier (canonical, then deprecated), THEN global-tier (canonical, then deprecated) + * -- the two tiers are never interleaved key-by-key, so this must stay two explicit bucket-tier + * lookups followed by two explicit global-tier lookups, not a single + * [[hadoopLookupPasswordEffective]] call per key (which would let an unset canonical bucket key + * fall through straight to the canonical GLOBAL value ahead of a SET deprecated bucket key, the + * wrong answer). + * + * The three-argument `lookupPassword` defaults to `""`, not `null`, so the last call in the + * block above never runs: Hadoop reads the deprecated global key only through `Configuration`'s + * deprecation alias of the canonical one. This function still reads it as a fourth lookup, + * which can only add a decline. + * + * THE FIX for the SSE-C long-bucket-alias gap is entirely inside the bucket tier: + * `lookupBucketSecret` itself is long-then-short, decompiled from `hadoop-aws` 3.3.4's + * `S3AUtils.class`: + * {{{ + * // longBucketKey = fs.s3a.bucket.B.fs.s3a. + * String longBucketKey = String.format(BUCKET_PATTERN, bucket, baseKey); + * String initialVal = getPassword(conf, longBucketKey, null, null); + * // shortBucketKey = fs.s3a.bucket.B. + * String shortBucketKey = String.format(BUCKET_PATTERN, bucket, subkey); + * // keeps initialVal (the LONG value) if non-empty + * return getPassword(conf, shortBucketKey, initialVal, null); + * }}} + * i.e. the SAME long-bucket-key construction and long-wins-if-nonempty semantics as + * `S3AUtils#lookupPassword` (see [[LookupPasswordConsumer]]/[[hadoopLookupPasswordEffective]]). + * `buildEncryptionSecrets` runs on the [[propagateBucketOptions]] view, so + * [[unsupportedEncryptionAlgorithmReason]] calls this on that view (and on the raw conf, which + * can only add a decline). An earlier version of this function modeled the bucket tier as + * SHORT-only, documented as "the LONG bucket form is genuinely never consulted for this key" -- + * that documentation was wrong (this decompilation supersedes it): a bucket configured only via + * `fs.s3a.bucket.B.fs.s3a.encryption.algorithm=SSE-C` bypassed the SSE-C gate entirely, because + * Hadoop's own reader DOES read that long form (and picks SSE-C), while this function reported + * `None` (nothing set) and the allowlist check below never even ran. + * + * The canonical-vs-deprecated distinction below is frequently moot in practice: `hadoop-aws`'s + * `S3AFileSystem.addDeprecatedKeys()` statically registers `fs.s3a.server-side-encryption-*` as + * `Configuration`-level deprecated aliases of `fs.s3a.encryption.*` (verified via `javap`), a + * registration that lives in a static field on Hadoop's `Configuration` class -- process-wide + * once `S3AFileSystem`'s class has loaded anywhere in the JVM, which a real scan has always + * already done by the time this gate runs, since reading the S3 table at all requires loading + * that class. Once active, `Configuration#get` resolves either literal key to the identical + * value transparently, making the two-key cascade below redundant (but harmless) for that case; + * it remains the operative path only when nothing else in the process has loaded + * `S3AFileSystem` yet. + * + * ALSO walks the Hadoop-credential-provider (JCEKS) path via [[resolveViaCredentialAliases]] + * for each of the four lookups below, matching `lookupBucketSecret`/`lookupPassword`'s real + * per-alias `getPassword` calls (quoted above) exactly: both are `getPassword`, not plain + * `Configuration#get`, so a bucket storing the algorithm name ONLY in a JCEKS keystore is + * exactly as real a Hadoop deployment shape for this key as it is for the credential keys + * [[hadoopLookupPasswordEffective]] already covers -- there is nothing algorithm-specific that + * makes JCEKS storage implausible here, so an earlier version of this function skipping it + * (documented at the time as "the algorithm NAME is not credential-sensitive data, so storing + * it in a keystore is not a realistic Hadoop deployment pattern") was an unjustified, narrower + * read than Hadoop's own resolver actually performs, under-declining a bucket whose algorithm + * is keystore-only. [[resolveViaCredentialAliases]]'s Arm B/C split still means this is zero + * extra I/O for the common case: keystore I/O only happens when a Hadoop credential-provider + * path is actually configured for the bucket, contained in that function's own try/catch. + * `bucketTier`/`globalTier` return `Left` (propagated straight through by [[orElseTier]]) when + * [[resolveViaCredentialAliases]] cannot safely verify a tier at all (an S3A-scoped provider + * path, or a corrupt/unreadable global keystore) -- correctly short-circuiting the whole + * cascade with a decline rather than silently falling through to a later tier that might look + * unset only because the true value was unverifiable. + */ + private def effectiveEncryptionAlgorithm( + hadoopConf: Configuration, + bucket: String): Either[String, Option[(String, String)]] = { + def bucketTier(baseKey: String): Either[String, Option[(String, String)]] = { + val longKey = s"fs.s3a.bucket.$bucket.$baseKey" + val shortKey = s"fs.s3a.bucket.$bucket." + baseKey.stripPrefix("fs.s3a.") + resolveViaCredentialAliases(hadoopConf, bucket, Seq(longKey, shortKey)) + .map(_.map(baseKey -> _)) + } + def globalTier(baseKey: String): Either[String, Option[(String, String)]] = + resolveViaCredentialAliases(hadoopConf, bucket, Seq(baseKey)) + .map(_.map(baseKey -> _)) + + // Short-circuits on Left (unverifiable tier) or Right(Some(_)) (resolved); only Right(None) + // (tier definitively unset) falls through to `next`, mirroring buildEncryptionSecrets's + // sequential `if (algorithm == null) algorithm = ...` cascade exactly. + def orElseTier( + current: Either[String, Option[(String, String)]], + next: => Either[String, Option[(String, String)]]) + : Either[String, Option[(String, String)]] = + current match { + case Left(reason) => Left(reason) + case Right(Some(value)) => Right(Some(value)) + case Right(None) => next + } + + orElseTier( + bucketTier(S3EncryptionAlgorithmKey), + orElseTier( + bucketTier(DeprecatedS3EncryptionAlgorithmKey), + orElseTier( + globalTier(S3EncryptionAlgorithmKey), + // Hadoop reaches this key only through the deprecation alias (see above); reading + // it directly can only over-decline. + globalTier(DeprecatedS3EncryptionAlgorithmKey)))) + } + + /** + * The only algorithm values a decline reason may print. Any other resolved value is withheld, + * since a `${...}` reference can expand the algorithm key to a customer key or token. + */ + private val PrintableEncryptionAlgorithms: Seq[String] = + Seq("AES256", "SSE-S3", "SSE-KMS", "DSSE-KMS", "SSE-C", "CSE-KMS", "CSE-CUSTOM") + + private def unsupportedEncryptionAlgorithmDeclineReason( + bucket: String, + algorithmKey: String, + algorithm: String): String = { + val setting = PrintableEncryptionAlgorithms + .find(_.equalsIgnoreCase(algorithm)) + .map(name => s"$algorithmKey=$name") + .getOrElse(s"$algorithmKey set to a value that is not a recognised S3 encryption algorithm") + s"Native Delta scan does not support $setting for $bucket " + + "(the native S3 client only supports unencrypted objects and S3's transparent " + + "server-side algorithms -- AES256/SSE-S3, SSE-KMS, and DSSE-KMS decrypt on GET/HEAD given " + + "read permission alone, with no extra request header; SSE-C additionally requires the " + + "customer-provided key resent as a header on every GET/HEAD request, which the native S3 " + + "client's extract_s3_config_options never forwards, and CSE-KMS/CSE-CUSTOM decrypt object " + + "bytes client-side, a layer the native Parquet reader does not have -- any of these would " + + "fail outright or silently read ciphertext where Hadoop's own reader succeeds)" + } + + /** + * First reason any bucket among `uris` is configured for an encryption algorithm the native S3 + * client cannot safely read, or `None` when claimable. Allowlist-based (see + * [[AllowedEncryptionAlgorithms]]): only `AES256`/`SSE-KMS`/`DSSE-KMS` (and unset/empty) pass; + * every other resolved value -- `SSE-C`, `CSE-KMS`, `CSE-CUSTOM`, or any unrecognized future + * algorithm string -- declines. Deliberately NOT a blocklist keyed on `SSE-C` alone: an + * allowlist is safe by construction against a Hadoop release adding a new encryption method + * this gate has never heard of, where a blocklist would silently admit it. Never interpolates a + * resolved key value, only key names, the bucket, and the (non-secret) algorithm name. Declines + * on a `Left` from [[effectiveEncryptionAlgorithm]] too (an unverifiable credential-provider + * arm, e.g. an S3A-scoped provider path or a corrupt/unreadable global keystore) -- the + * algorithm cannot be ruled safe when it cannot be read at all. + */ + private[delta] def unsupportedEncryptionAlgorithmReason( + hadoopConf: Configuration, + uris: Seq[URI], + propagatedConfCache: MutableMap[String, Configuration] = MutableMap.empty) + : Option[String] = { + val buckets = uris.flatMap(s3Bucket).distinct + buckets.foldLeft(Option.empty[String]) { (declined, bucket) => + if (declined.isDefined) { + declined + } else { + try { + // Hadoop resolves the algorithm on the propagated view, where a bucket override can + // change what a `${...}` reference expands to. The raw view is checked as well, so a + // disagreement between the two can only decline. + val propagatedConf = + propagatedConfCache.getOrElseUpdate( + bucket, + propagateBucketOptions(hadoopConf, bucket)) + encryptionAlgorithmReason(propagatedConf, bucket) + .orElse(encryptionAlgorithmReason(hadoopConf, bucket)) + } catch { + case e @ (_: IOException | _: RuntimeException) => + Some(unverifiableValueReason(bucket, S3EncryptionAlgorithmKey, e)) + } + } + } + } + + private def encryptionAlgorithmReason(conf: Configuration, bucket: String): Option[String] = + effectiveEncryptionAlgorithm(conf, bucket) match { + case Left(reason) => Some(reason) + case Right(None) => None + case Right(Some((key, value))) => + if (!AllowedEncryptionAlgorithms.exists(_.equalsIgnoreCase(value))) { + Some(unsupportedEncryptionAlgorithmDeclineReason(bucket, key, value)) + } else { + None + } + } + + /** + * Canonical Hadoop S3A HTTP-proxy host config key, verified via CFR decompilation of + * `hadoop-aws` 3.3.4's `S3AUtils.class` (`initProxySupport`): + * {{{ + * String proxyHost = conf.getTrimmed("fs.s3a.proxy.host", ""); + * int proxyPort = conf.getInt("fs.s3a.proxy.port", -1); + * if (!proxyHost.isEmpty()) { + * ... + * String proxyUsername = + * S3AUtils.lookupPassword(bucket, conf, "fs.s3a.proxy.username", null, null); + * String proxyPassword = + * S3AUtils.lookupPassword(bucket, conf, "fs.s3a.proxy.password", null, null); + * ... + * } + * }}} + * `fs.s3a.proxy.host`/`fs.s3a.proxy.port` resolve via a PLAIN, non-bucket-scoped, non-JCEKS + * `Configuration#getTrimmed`/`getInt` call -- NOT `lookupPassword` -- against whatever conf + * `S3AFileSystem#initialize` already ran through `propagateBucketOptions` before + * `createAwsConf`/`initProxySupport` ever runs; only the SIBLING `fs.s3a.proxy.username`/ + * `fs.s3a.proxy.password` keys go through `lookupPassword` (bucket long/short/global, + * JCEKS-aware). So the host is bucket-aware only via `propagateBucketOptions`'s short-bucket- + * form folding, never the long-bucket form, and never a credential-provider read -- the SAME + * shape as `endpoint`/`path.style.access` ([[PropagatedOptionConsumer]], see + * [[S3ConfigKeyConsumers]]'s doc), not the credential family. `hadoop-aws` 3.4.x (the Spark 4.x + * profiles' version) moves this code to `AWSClientConfig#createProxyConfiguration`/ + * `#createAsyncProxyConfiguration` but keeps the exact same reads, verified via `javap` against + * 3.4.2: `conf.getTrimmed("fs.s3a.proxy.host", "")` for the host, `S3AUtils.lookupPassword` for + * username/password only. + * + * [[proxyGateReason]] therefore resolves the host EXACTLY like its real consumer -- plain + * `Configuration#getTrimmed` on the [[propagateBucketOptions]] result, no provider arms, no + * [[ClearTextFallbackKey]] handling -- rather than through the wider `lookupPassword` cascade + * an earlier version reused here. The wider cascade was wrong in BOTH directions for this key: + * with only the GLOBAL Hadoop provider path set and [[ClearTextFallbackKey]] false, + * `getPassword` hides a plaintext host that `getTrimmed` serves to Hadoop anyway (a missed + * decline, the exact bypass this gate exists to close -- native has NO HTTP-proxy support of + * any kind, no `fs.s3a.proxy.*` key is read anywhere in `s3.rs`); and an S3A-scoped provider + * path or a lone long-form bucket alias declined a bucket whose real consumer can never see a + * host from either source (pure over-refusal -- no keystore can supply the host to a plain + * `getTrimmed`, and the long form folds into the unread `fs.s3a.fs.s3a.proxy.host`). + */ + private val S3ProxyHostKey = "fs.s3a.proxy.host" + + private def unsupportedProxyReason(bucket: String, key: String): String = + s"Native Delta scan does not support $key configured for $bucket (the native S3 client has " + + "no HTTP proxy support at all -- no fs.s3a.proxy.* key is read anywhere in its object " + + "store layer -- so a claimed scan would connect to S3 directly instead of routing through " + + "the configured proxy, either bypassing an egress/network-segmentation policy or simply " + + "failing to reach the endpoint)" + + /** + * First reason any bucket among `uris` has an HTTP proxy configured via [[S3ProxyHostKey]], or + * `None` when claimable. Reads the host EXACTLY like its real consumer (see + * [[S3ProxyHostKey]]'s doc): plain `Configuration#getTrimmed` on the [[propagateBucketOptions]] + * result, so this gate is always zero-I/O -- no credential provider is ever consulted for the + * host, because none ever supplies it to Hadoop either. Ordered alongside + * [[unsupportedEncryptionAlgorithmReason]] among the conf-only gates, ahead of + * [[s3ConfigDivergenceReason]]. The try/catch guards `Configuration#get`'s + * `IllegalStateException` on a `${...}` substitution cycle, same as [[s3KeyDivergenceReason]]. + * Never interpolates a resolved value: the proxy HOST is not secret, but naming it here would + * be a strange place to first surface it, and proxy CREDENTIALS + * (`fs.s3a.proxy.username`/`fs.s3a.proxy.password`, not read by this gate at all -- the whole + * point of gating on the host is that a non-empty host declines before any proxy credential + * would ever need to be forwarded) must never appear in a decline reason regardless. + */ + private[delta] def proxyGateReason( + hadoopConf: Configuration, + uris: Seq[URI], + propagatedConfCache: MutableMap[String, Configuration] = MutableMap.empty) + : Option[String] = { + val buckets = uris.flatMap(s3Bucket).distinct + buckets.foldLeft(Option.empty[String]) { (declined, bucket) => + if (declined.isDefined) { + declined + } else { + try { + val propagatedConf = + propagatedConfCache.getOrElseUpdate( + bucket, + propagateBucketOptions(hadoopConf, bucket)) + if (propagatedConf.getTrimmed(S3ProxyHostKey, "").nonEmpty) { + Some(unsupportedProxyReason(bucket, S3ProxyHostKey)) + } else { + None + } + } catch { + case e @ (_: IOException | _: RuntimeException) => + Some(unverifiableValueReason(bucket, S3ProxyHostKey, e)) + } + } + } + } + + /** + * Canonical Hadoop S3A assumed-role session-policy key. `AssumedRoleCredentialProvider`'s + * constructor reads it via a plain `Configuration#getTrimmed` against the + * [[propagateBucketOptions]]-propagated conf (verified via `javap` against `hadoop-aws` 3.3.4 + * and 3.4.1: `conf.getTrimmed("fs.s3a.assumed.role.policy", "")` -- the same consumer shape as + * `assumed.role.arn`/`session.name`, see [[S3ConfigKeyConsumers]]) and, when non-empty, + * attaches it as the session policy of its STS AssumeRole request. + */ + private val S3AssumedRolePolicyKey = "fs.s3a.assumed.role.policy" + + private def assumedRolePolicyReason(bucket: String, key: String): String = + s"Native Delta scan does not support $key configured for $bucket (Hadoop sends the " + + "configured session policy in its STS AssumeRole request, but the native S3 client's " + + "assumed-role provider never reads or forwards this key, so a claimed scan would " + + "assume the role WITHOUT the configured session restriction -- silently widening the " + + "effective permissions instead of failing)" + + /** + * First reason any bucket among `uris` configures an assumed-role session policy via + * [[S3AssumedRolePolicyKey]], or `None` when claimable. Hadoop's + * `AssumedRoleCredentialProvider` includes the policy in its AssumeRole request; native's + * assumed-role provider does not, so a configured policy must decline until native supports it. + * Resolved exactly like the key's real consumer (plain `getTrimmed` on the propagated conf, + * mirroring [[proxyGateReason]]); declined whenever set, whether or not the current provider + * chain names the assumed-role provider -- a policy that is dead config today can become live + * through a provider-chain change native never re-validates. Never interpolates the policy + * document itself, only the key and bucket. + */ + private[delta] def assumedRolePolicyGateReason( + hadoopConf: Configuration, + uris: Seq[URI], + propagatedConfCache: MutableMap[String, Configuration] = MutableMap.empty) + : Option[String] = { + val buckets = uris.flatMap(s3Bucket).distinct + buckets.foldLeft(Option.empty[String]) { (declined, bucket) => + if (declined.isDefined) { + declined + } else { + try { + val propagatedConf = + propagatedConfCache.getOrElseUpdate( + bucket, + propagateBucketOptions(hadoopConf, bucket)) + if (propagatedConf.getTrimmed(S3AssumedRolePolicyKey, "").nonEmpty) { + Some(assumedRolePolicyReason(bucket, S3AssumedRolePolicyKey)) + } else { + None + } + } catch { + case e @ (_: IOException | _: RuntimeException) => + Some(unverifiableValueReason(bucket, S3AssumedRolePolicyKey, e)) + } + } + } + } + + /** Hadoop's S3A SSL switch; `S3AFileSystem` reads it via `Configuration#getBoolean`. */ + private val S3SslEnabledKey = "fs.s3a.connection.ssl.enabled" + private val S3EndpointKey = "fs.s3a.endpoint" + + /** + * Hadoop's assumed-role STS endpoint keys; `AssumedRoleCredentialProvider` sends AssumeRole to + * the configured endpoint, while native builds its provider with SDK defaults. + */ + private val S3StsEndpointKeys = + Seq("fs.s3a.assumed.role.sts.endpoint", "fs.s3a.assumed.role.sts.endpoint.region") + + private def insecureEndpointReason(bucket: String): String = + s"Native Delta scan does not support a scheme-less $S3EndpointKey with " + + s"$S3SslEnabledKey=false for $bucket (Hadoop addresses that endpoint over http://, " + + "while the native S3 client always assumes https://, so a claimed scan would fail at " + + "execution where Spark reads fine)" + + private def stsEndpointReason(bucket: String, key: String): String = + s"Native Delta scan does not support $key configured for $bucket (Hadoop sends its " + + "AssumeRole request to the configured STS endpoint, while the native S3 client's " + + "assumed-role provider uses the SDK defaults, so the two sides would authenticate " + + "against different endpoints)" + + /** + * First reason any bucket among `uris` configures an endpoint native would address differently + * from Hadoop: a scheme-less `fs.s3a.endpoint` with SSL disabled, or an assumed-role STS + * endpoint. Resolved like the keys' real consumers (plain reads on the propagated conf, + * mirroring [[proxyGateReason]]); the STS keys decline whenever set, since dead config can + * become live through a provider-chain change native never re-validates. + */ + private[delta] def hadoopOnlyEndpointGateReason( + hadoopConf: Configuration, + uris: Seq[URI], + propagatedConfCache: MutableMap[String, Configuration] = MutableMap.empty) + : Option[String] = { + val buckets = uris.flatMap(s3Bucket).distinct + buckets.foldLeft(Option.empty[String]) { (declined, bucket) => + if (declined.isDefined) { + declined + } else { + try { + val propagatedConf = + propagatedConfCache.getOrElseUpdate( + bucket, + propagateBucketOptions(hadoopConf, bucket)) + val endpoint = propagatedConf.getTrimmed(S3EndpointKey, "") + val schemeless = endpoint.nonEmpty && !endpoint.contains("://") + if (schemeless && !propagatedConf.getBoolean(S3SslEnabledKey, true)) { + Some(insecureEndpointReason(bucket)) + } else { + S3StsEndpointKeys + .find(key => propagatedConf.getTrimmed(key, "").nonEmpty) + .map(key => stsEndpointReason(bucket, key)) + } + } catch { + case e @ (_: IOException | _: RuntimeException) => + Some(unverifiableValueReason(bucket, S3EndpointKey, e)) + } + } + } + } + + /** + * `baseKey`'s long, short, then global per-bucket aliases, in Hadoop's own resolution order. + */ + private def longThenShortThenGlobalAliases(bucket: String, baseKey: String): Seq[String] = { + val suffix = baseKey.stripPrefix("fs.s3a.") + Seq(s"fs.s3a.bucket.$bucket.fs.s3a.$suffix", s"fs.s3a.bucket.$bucket.$suffix", baseKey) + } + + private def s3aScopedProviderPathReason(bucket: String, providerPathKey: String): String = + "Native Delta scan cannot forward Hadoop credential-provider aliases for " + + s"$bucket ($providerPathKey configures an S3A-scoped Hadoop credential provider that " + + "Configuration#getPassword does not consult, so the native S3 client's credentials " + + "cannot be verified)" + + private def unverifiableCredentialProviderReason(bucket: String, error: Throwable): String = + "Native Delta scan cannot verify Hadoop credential-provider aliases for " + + s"$bucket (reading $HadoopCredentialProviderPathKey raised " + + s"${error.getClass.getName}), declining rather than risk missing credentials" + + /** + * Three-way Hadoop credential-provider-path precheck shared by every `getPassword`-based + * resolution below, a property of the bucket alone. An S3A- or bucket-scoped provider path, + * which `Configuration#getPassword` never consults, yields [[UnverifiableProvider]] with no + * keystore I/O. A global path with any entry that is not a local keystore yields + * [[NonLocalGlobalProvider]], also with no I/O (see [[globalProviderPathHasNonLocalEntry]]). + * Only a global path of local entries yields [[GlobalProviderOnly]], the one arm whose callers + * do real keystore I/O and must wrap their own `getPassword` calls in try/catch. No provider + * path anywhere yields [[NoProvider]], zero-I/O plain conf reads only. + */ + private sealed trait CredentialProviderArm + private case class UnverifiableProvider(offendingKey: String) extends CredentialProviderArm + private case object NonLocalGlobalProvider extends CredentialProviderArm + private case object GlobalProviderOnly extends CredentialProviderArm + private case object NoProvider extends CredentialProviderArm + + /** + * True when the global provider path has an entry other than a local keystore or a + * filesystem-free provider. An entry passes only as `jceks`, `localjceks`, `bcfks` or + * `localbcfks` with nested scheme `file` (as in `jceks://file/...`) while `fs.file.impl` is + * unset or `org.apache.hadoop.fs.LocalFileSystem`, or as `user` with no authority (`user:///`). + * Any other entry would run a third-party provider or open a filesystem during planning, and + * S3AFileSystem drops entries whose filesystem class is S3A-compatible before it resolves + * anything (`ProviderUtils.excludeIncompatibleCredentialProviders`), so reading one could also + * disagree with Hadoop. An unparseable or empty entry counts as non-local. The nested scheme is + * the part of the authority before `@`, as in `ProviderUtils.unnestUri`. + */ + private def globalProviderPathHasNonLocalEntry(hadoopConf: Configuration): Boolean = + Option(hadoopConf.get(HadoopCredentialProviderPathKey)).filter(_.nonEmpty).exists { path => + val isLocalFileSystem = Option(hadoopConf.get("fs.file.impl")) + .forall(_ == "org.apache.hadoop.fs.LocalFileSystem") + path.split(",").exists(entry => !isLocalOrFilesystemFreeProvider(entry, isLocalFileSystem)) + } + + private val KeystoreProviderSchemes = Set("jceks", "localjceks", "bcfks", "localbcfks") + + private def isLocalOrFilesystemFreeProvider( + entry: String, + isLocalFileSystem: Boolean): Boolean = + try { + val uri = new URI(entry) + Option(uri.getAuthority) match { + case Some(authority) => + KeystoreProviderSchemes.contains(uri.getScheme) && isLocalFileSystem && + authority.split("@", 2)(0) == "file" + case None => uri.getScheme == "user" + } + } catch { + case _: URISyntaxException => false + } + + private def nonLocalProviderPathReason(bucket: String): String = + s"Native Delta scan cannot verify the S3 settings for $bucket " + + s"($HadoopCredentialProviderPathKey names a credential keystore that is not on the local " + + "filesystem, and the settings cannot be verified without reading it)" + + private def credentialProviderArm( + hadoopConf: Configuration, + bucket: String): CredentialProviderArm = { + val bucketPathKey = s3aBucketProviderPathKey(bucket) + val bucketLongPathKey = s3aBucketLongProviderPathKey(bucket) + val s3aPathSet = nonEmptyConf(hadoopConf, S3aCredentialProviderPathKey) + val bucketPathSet = nonEmptyConf(hadoopConf, bucketPathKey) + val bucketLongPathSet = nonEmptyConf(hadoopConf, bucketLongPathKey) + if (s3aPathSet || bucketPathSet || bucketLongPathSet) { + val offendingKey = + if (s3aPathSet) S3aCredentialProviderPathKey + else if (bucketPathSet) bucketPathKey + else bucketLongPathKey + UnverifiableProvider(offendingKey) + } else if (globalProviderPathHasNonLocalEntry(hadoopConf)) { + NonLocalGlobalProvider + } else if (nonEmptyConf(hadoopConf, HadoopCredentialProviderPathKey)) { + GlobalProviderOnly + } else { + NoProvider + } + } + + /** + * Resolves `aliases` in order under `hadoopConf`/`bucket`, keeping the first non-empty value, + * or `Left(reason)` when the value cannot be safely verified. Dispatches on + * [[credentialProviderArm]]: [[UnverifiableProvider]] and [[NonLocalGlobalProvider]] decline + * with zero I/O; [[GlobalProviderOnly]] resolves each alias via `Configuration#getPassword`, + * keystore I/O contained in try/catch so a corrupt store declines this bucket rather than + * aborting planning; [[NoProvider]] resolves each alias via zero-I/O [[plainValue]] reads, + * which honor [[ClearTextFallbackKey]] the way a real `getPassword` consumer would. Callers + * such as [[hadoopLookupPasswordEffective]] and [[effectiveEncryptionAlgorithm]] supply the + * alias lists that match their consumer's real per-tier `getPassword` calls. + */ + private def resolveViaCredentialAliases( + hadoopConf: Configuration, + bucket: String, + aliases: Seq[String]): Either[String, Option[String]] = + credentialProviderArm(hadoopConf, bucket) match { + case UnverifiableProvider(offendingKey) => + Left(s3aScopedProviderPathReason(bucket, offendingKey)) + case NonLocalGlobalProvider => + Left(nonLocalProviderPathReason(bucket)) + case GlobalProviderOnly => + try { + Right( + aliases.iterator + .map(alias => + Option(hadoopConf.getPassword(alias)).map(new String(_)).filter(_.nonEmpty)) + .collectFirst { case Some(v) => v }) + } catch { + case e @ (_: IOException | _: RuntimeException) => + Left(unverifiableCredentialProviderReason(bucket, e)) + } + case NoProvider => + if (!hadoopConf.getBoolean(ClearTextFallbackKey, true)) { + Right(None) + } else { + Right(aliases.flatMap(plainValue(hadoopConf, _)).headOption) + } + } + + /** + * `bucket`'s effective value for `baseKey` in Hadoop's own `S3AUtils#lookupPassword` resolution + * order -- long bucket alias, then short bucket alias, then global, each tried through a Hadoop + * credential provider before falling back to plain conf (see [[resolveViaCredentialAliases]] + * for the Arm A/B/C dispatch this delegates to) -- or `Left(reason)` when the value cannot be + * safely verified. + * + * USED ONLY for keys whose real consumer IS `lookupPassword` -- the [[LookupPasswordConsumer]] + * entries of [[S3ConfigKeyConsumers]], via [[s3KeyDivergenceReason]]. An earlier version ran + * EVERY compared key through this function, reasoning that a wider read could only ever + * over-decline; for a value-EQUALITY comparator that reasoning is half-true: the long-form + * alias this function consults FIRST can hold exactly the value native resolves while Hadoop's + * true propagate-then-plain-get value differs (e.g. a `${...}` reference whose referent + * propagation redirects), producing a false EQUALITY that admits a really-diverging scan. A + * [[PropagatedOptionConsumer]] key must therefore resolve like its actual consumer instead -- + * see [[s3KeyDivergenceReason]]. + * + * Honors [[ClearTextFallbackKey]] (via [[resolveViaCredentialAliases]]'s [[NoProvider]] arm), + * matching `getPassword`'s real refusal to read plaintext conf when the flag is off. + * + * `hadoopConf` must be a [[propagateBucketOptions]] result (the caller, + * [[s3KeyDivergenceReason]], always passes one) so that `${...}` references embedded in any + * alias resolve exactly like `S3AFileSystem#initialize`'s real propagate-then-resolve order. + */ + private def hadoopLookupPasswordEffective( + hadoopConf: Configuration, + bucket: String, + baseKey: String): Either[String, Option[String]] = + resolveViaCredentialAliases( + hadoopConf, + bucket, + longThenShortThenGlobalAliases(bucket, baseKey)) + + private def effectiveValueDivergenceReason(bucket: String, key: String): String = + s"Native Delta scan cannot forward $key for $bucket (Hadoop's effective value for this key " + + "differs from what the native S3 client resolves, so its credentials or configuration " + + "would differ from Hadoop's)" + + private def unverifiableValueReason(bucket: String, key: String, error: Throwable): String = + s"Native Delta scan cannot verify $key for $bucket (Configuration#get raised " + + s"${error.getClass.getName}), declining rather than risk forwarding a stale or " + + "diverging value" + + /** + * `None` when `baseKey`'s Hadoop-effective and native-effective values under `bucket` agree, or + * a decline reason naming `baseKey` and `bucket` (never a value) when they diverge or either + * side cannot be safely computed. + * + * Hadoop's effective value is computed against [[propagateBucketOptions]]'s result, mirroring + * `S3AFileSystem#initialize`'s actual order (propagate bucket options into the conf FIRST, only + * THEN read/substitute options against it), through the resolution `consumer` declares for the + * key in [[S3ConfigKeyConsumers]]: [[LookupPasswordConsumer]] keys via + * [[hadoopLookupPasswordEffective]] (long-then-short-then-global, keystore- and + * [[ClearTextFallbackKey]]-aware), [[PropagatedOptionConsumer]] keys via a plain + * `Configuration#get` on the propagated view (short-form wins by propagation alone; the long + * form and any keystore/fallback handling are ignored, exactly like the key's real consumer). + * Resolving a plain-consumer key through the wider `lookupPassword` cascade instead would let a + * long-form alias value Hadoop never reads EQUAL native's resolution while Hadoop's true + * plain-get value differs -- a false equality admitting a diverging scan, not merely an extra + * decline. Native's effective value is always `nativeShortThenGlobal(hadoopConf, ...)` on the + * ORIGINAL, unpropagated conf, matching `NativeConfig.extractObjectStoreOptions`'s actual + * forwarding semantics (no propagation step) AND native's `get_config` presence-based (not + * emptiness-based) short-vs-global fallback -- see [[nativeShortThenGlobal]]. Both values are + * trimmed together, symmetrically, right before the equality check below (mirroring + * `get_config_trimmed`'s `.trim()`, which native applies regardless of which alias it read) + * rather than trimming [[nativeShortThenGlobal]]'s result on its own -- see + * [[nativeShortThenGlobal]]'s doc for why a one-sided trim there would flag a spurious + * divergence. Comparing against the propagated view (rather than the original conf, as an + * earlier version of this check did) matters because propagation can change what a `${...}` + * reference inside one bucket-scoped value resolves to: e.g. + * `fs.s3a.bucket.B.access.key=${fs.s3a.custom.ref}` with `fs.s3a.bucket.B.custom.ref=X` and + * global `fs.s3a.custom.ref=Y` propagates to `fs.s3a.custom.ref=X` (overwriting the global `Y`) + * before the access key's `${...}` reference is ever substituted, so Hadoop resolves `X` while + * a check against the unpropagated conf would (wrongly) also see `Y`, the same value native + * forwards -- masking a real divergence. Wrapped in try/catch: `Configuration#get` raises + * `IllegalStateException` once `${...}` substitution recurses past Hadoop's `MAX_SUBST` bound + * (e.g. a two-key mutual reference cycle); declining is safer than crashing planning or + * comparing a partially-substituted value. + */ + private def s3KeyDivergenceReason( + hadoopConf: Configuration, + propagatedConf: Configuration, + bucket: String, + baseKey: String, + consumer: S3ConfigConsumer): Option[String] = { + try { + val hadoopEffective: Either[String, Option[String]] = consumer match { + case LookupPasswordConsumer => + hadoopLookupPasswordEffective(propagatedConf, bucket, baseKey) + case PropagatedOptionConsumer => + Right(Option(propagatedConf.get(baseKey))) + } + hadoopEffective match { + case Left(reason) => Some(reason) + case Right(hadoopValue) => + val nativeValue = nativeShortThenGlobal(hadoopConf, bucket, baseKey) + if (hadoopValue.map(_.trim) != nativeValue.map(_.trim)) { + Some(effectiveValueDivergenceReason(bucket, baseKey)) + } else { + None + } + } + } catch { + case e @ (_: IOException | _: RuntimeException) => + Some(unverifiableValueReason(bucket, baseKey, e)) + } + } + + /** + * First reason any bucket among `uris` cannot faithfully forward every [[AllS3ConfigKeys]] + * option to native, or `None` when every key's Hadoop-effective and native-effective value + * agrees for every S3/S3A bucket referenced. Only `s3`/`s3a` authorities matter here (ABFS/WASB + * mooted by the userinfo gate, GCS handled by [[gcsHadoopOnlyAuthReason]]). One comparator + * replaces the former per-case gate family (long-form bucket credentials, JCEKS/provider + * shadowing, Hadoop `${...}` variable references): [[s3KeyDivergenceReason]] computes Hadoop's + * effective value against a per-bucket [[propagateBucketOptions]] replica (matching + * `S3AFileSystem#initialize`'s real propagate-then-resolve order), so a `${...}` reference that + * resolves identically under that propagated view and under native's unpropagated forwarding is + * no longer a divergence at all, while one that resolves differently (e.g. because propagation + * shadowed a referenced key with a per-bucket override) IS still caught. + * [[s3KeyDivergenceReason]] resolves each key through the consumer family + * [[S3ConfigKeyConsumers]] declares for it, right beside the key itself: the SSE-C + * long-bucket-alias bypass came from a key's resolution being decided implicitly, scattered + * across call sites, so the classification is now a single visible list -- and the tier must + * MATCH the key's real consumer in both directions, because an equality comparator resolving a + * plain-get key through the wider `lookupPassword` cascade can manufacture a false EQUALITY + * (long-form alias equal to native's value, true propagated plain value different) just as + * readily as a false divergence. Never interpolates a resolved value, only key names. + */ + private[delta] def s3ConfigDivergenceReason( + hadoopConf: Configuration, + uris: Seq[URI], + propagatedConfCache: MutableMap[String, Configuration] = MutableMap.empty) + : Option[String] = { + val buckets = uris.flatMap(s3Bucket).distinct + buckets.foldLeft(Option.empty[String]) { (declined, bucket) => + if (declined.isDefined) { + declined + } else { + try { + val propagatedConf = + propagatedConfCache.getOrElseUpdate( + bucket, + propagateBucketOptions(hadoopConf, bucket)) + S3ConfigKeyConsumers.foldLeft(Option.empty[String]) { + case (keyDeclined, (key, consumer)) => + if (keyDeclined.isDefined) { + keyDeclined + } else { + s3KeyDivergenceReason(hadoopConf, propagatedConf, bucket, key, consumer) + } + } + } catch { + case e @ (_: IOException | _: RuntimeException) => + Some(unverifiableValueReason(bucket, AllS3ConfigKeys.head, e)) + } + } + } + } + + /** + * String-literal mirror of every credential-provider class name s3.rs's + * `build_aws_credential_provider_metadata` recognizes (Hadoop S3A plus AWS SDK v1/v2 names). + * `hadoop-aws` is NOT on this module's runtime classpath, so these stay string literals, never + * `classOf` references. + */ + private val SupportedCredentialProviderClasses: Set[String] = Set( + "org.apache.hadoop.fs.s3a.auth.IAMInstanceCredentialsProvider", + "org.apache.hadoop.fs.s3a.SimpleAWSCredentialsProvider", + "org.apache.hadoop.fs.s3a.TemporaryAWSCredentialsProvider", + "org.apache.hadoop.fs.s3a.auth.AssumedRoleCredentialProvider", + "org.apache.hadoop.fs.s3a.AnonymousAWSCredentialsProvider", + "software.amazon.awssdk.auth.credentials.ContainerCredentialsProvider", + "com.amazonaws.auth.ContainerCredentialsProvider", + "com.amazonaws.auth.EC2ContainerCredentialsProviderWrapper", + "software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider", + "com.amazonaws.auth.InstanceProfileCredentialsProvider", + "software.amazon.awssdk.auth.credentials.EnvironmentVariableCredentialsProvider", + "com.amazonaws.auth.EnvironmentVariableCredentialsProvider", + "software.amazon.awssdk.auth.credentials.WebIdentityTokenFileCredentialsProvider", + "com.amazonaws.auth.WebIdentityTokenCredentialsProvider", + "software.amazon.awssdk.auth.credentials.ProfileCredentialsProvider", + "com.amazonaws.auth.profile.ProfileCredentialsProvider", + "software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider", + "com.amazonaws.auth.AnonymousAWSCredentials") + + private val AnonymousCredentialProviderClasses: Set[String] = Set( + "org.apache.hadoop.fs.s3a.AnonymousAWSCredentialsProvider", + "software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider", + "com.amazonaws.auth.AnonymousAWSCredentials") + + private val HadoopAssumedRoleProviderClass = + "org.apache.hadoop.fs.s3a.auth.AssumedRoleCredentialProvider" + + private val AwsCredentialsProviderKey = "fs.s3a.aws.credentials.provider" + private val AssumedRoleCredentialsProviderKey = "fs.s3a.assumed.role.credentials.provider" + + /** Splits a comma-separated credential-provider-class list the same way s3.rs's parser does. */ + private def parseProviderClassNames(value: String): Seq[String] = + value.split(",").map(_.trim).filter(_.nonEmpty).toSeq + + // The class name is withheld: an unsupported value is never one of the known classes, and a + // `${...}` reference can expand the key to a credential. + private def unsupportedProviderReason(bucket: String, key: String): String = + "Native Delta scan does not support a credential provider class " + + s"configured via $key for $bucket (the native S3 client only supports a fixed set of " + + "provider classes; an unsupported class would fail at scan execution time, after the " + + "scan was already claimed, rather than at planning time)" + + private def mixedAnonymousProviderReason(bucket: String, key: String): String = + s"Native Delta scan does not support $key for $bucket naming an anonymous credential " + + "provider together with any other provider (the native S3 client rejects this " + + "combination at scan execution time)" + + private def anonymousAssumedRoleProviderReason(bucket: String, key: String): String = + s"Native Delta scan does not support an anonymous credential provider in $key for " + + s"$bucket (the native S3 client does not allow an anonymous provider as the base " + + "credentials for an assumed-role chain)" + + private def unsupportedProviderNameReason( + bucket: String, + key: String, + names: Seq[String]): Option[String] = + if (names.exists(name => !SupportedCredentialProviderClasses.contains(name))) { + Some(unsupportedProviderReason(bucket, key)) + } else { + None + } + + /** + * [[shortThenGlobal]] for `key` under `bucket`, or `Left(reason)` when `Configuration#get` + * itself raises: Hadoop throws `IllegalStateException` once `${...}` expansion recurses past + * its `MAX_SUBST` bound, e.g. a mutual reference cycle between two provider keys. Every + * provider-class read below goes through this wrapper so the exception is caught here, + * whichever entry point runs first. Reuses [[unverifiableValueReason]]'s message shape (names + * the key and the exception class, never a value). + */ + private def shortThenGlobalOrReason( + hadoopConf: Configuration, + bucket: String, + key: String): Either[String, Option[String]] = + try { + Right(shortThenGlobal(hadoopConf, bucket, key)) + } catch { + case e @ (_: IOException | _: RuntimeException) => + Left(unverifiableValueReason(bucket, key, e)) + } + + /** + * Decline reason when `bucket`'s effective `assumed.role.credentials.provider` names an + * unsupported class, or an anonymous one (native rejects ANY anonymous entry here, not just a + * mix), or when reading it raises (see [[shortThenGlobalOrReason]]). Unset defaults to native's + * own always-supported fallback, so `None` is safe. + */ + private def assumedRoleProviderClassReason( + hadoopConf: Configuration, + bucket: String): Option[String] = { + shortThenGlobalOrReason(hadoopConf, bucket, AssumedRoleCredentialsProviderKey) match { + case Left(reason) => Some(reason) + case Right(None) => None + case Right(Some(value)) => + val names = parseProviderClassNames(value) + unsupportedProviderNameReason(bucket, AssumedRoleCredentialsProviderKey, names).orElse { + if (names.exists(AnonymousCredentialProviderClasses.contains)) { + Some(anonymousAssumedRoleProviderReason(bucket, AssumedRoleCredentialsProviderKey)) + } else { + None + } + } + } + } + + /** + * Decline reason when `bucket`'s effective `aws.credentials.provider` names an unrecognized + * class, mixes an anonymous provider with any other, a nested `AssumedRoleCredentialProvider` + * sub-chain has the same problem, or reading either key raises (see + * [[shortThenGlobalOrReason]]). Unset/empty falls back to native's default chain. + */ + private def providerClassReason(hadoopConf: Configuration, bucket: String): Option[String] = { + shortThenGlobalOrReason(hadoopConf, bucket, AwsCredentialsProviderKey) match { + case Left(reason) => Some(reason) + case Right(None) => None + case Right(Some(value)) => + val names = parseProviderClassNames(value) + unsupportedProviderNameReason(bucket, AwsCredentialsProviderKey, names) + .orElse { + if (names.length > 1 && names.exists(AnonymousCredentialProviderClasses.contains)) { + Some(mixedAnonymousProviderReason(bucket, AwsCredentialsProviderKey)) + } else { + None + } + } + .orElse { + if (names.contains(HadoopAssumedRoleProviderClass)) { + assumedRoleProviderClassReason(hadoopConf, bucket) + } else { + None + } + } + } + } + + /** + * First reason any bucket among `uris` names an unsupported credential-provider class, or + * `None` when every named class is supported (or the key is unset). `NativeConfig` forwards + * `Configuration#get`'s substituted value for every entry, the same read [[shortThenGlobal]] + * performs here, so a `${...}` reference resolves identically for native and for this check. + */ + private[delta] def providerClassGateReason( + hadoopConf: Configuration, + uris: Seq[URI]): Option[String] = { + val buckets = uris.flatMap(s3Bucket).distinct + buckets.foldLeft(Option.empty[String]) { (declined, bucket) => + if (declined.isDefined) declined else providerClassReason(hadoopConf, bucket) + } + } + + /** + * True when `key` names a GCS authentication option under either Hadoop conf namespace the + * `gcs-connector` reads (`fs.gs.*` or the legacy `google.cloud.*`) AND the key itself concerns + * authentication. The connector's own `HadoopCredentialConfiguration` builds each auth setting + * from a prefix crossed with a suffix (service-account keyfile/email/private-key, OAuth client + * id/secret, impersonation, workload identity, and so on), including reversed-word-order + * deprecated forms (`fs.gs.service.account.auth.keyfile`) alongside the modern ones + * (`fs.gs.auth.service.account.json.keyfile`) -- enumerating every current and future suffix as + * a fixed prefix list is a losing game the connector itself does not play; matching on + * "namespace + contains auth" tracks the connector's own auth-vs-non-auth boundary instead of + * chasing its naming history. `gcs-connector` is NOT on this module's runtime classpath by + * default, so referencing an actual GCS auth class would risk `NoClassDefFoundError`, same + * rationale as the S3A literals above. + */ + private def isGcsAuthKey(key: String): Boolean = + (key.startsWith("fs.gs.") || key.startsWith("google.cloud.")) && key.contains("auth") + + /** + * True when `uri`'s scheme is `gs` (case-insensitive) -- the ONLY scheme object_store's + * `ObjectStoreScheme::parse` (parquet_support.rs) routes to `GoogleCloudStorage`; `gcs` is not + * recognized there and is deliberately excluded. + */ + private def isGcsScheme(uri: URI): Boolean = + Option(uri.getScheme).exists(_.equalsIgnoreCase("gs")) + + /** + * The lowercase-scheme-checked GCS bucket name from `uri`'s authority (host, minus any userinfo + * or port), or `None` when `uri`'s scheme is not `gs`. Parses the raw authority manually, + * mirroring [[s3Bucket]]'s `URI#getHost`/RFC 3986 `reg-name` reasoning. + */ + private def gcsBucket(uri: URI): Option[String] = { + if (!isGcsScheme(uri)) { + None + } else { + val authority = Option(uri.getAuthority).getOrElse("") + val at = authority.lastIndexOf('@') + val hostAndPort = if (at >= 0) authority.substring(at + 1) else authority + val colon = hostAndPort.lastIndexOf(':') + val host = if (colon >= 0) hostAndPort.substring(0, colon) else hostAndPort + if (host.isEmpty) None else Some(host) + } + } + + /** + * The non-empty Hadoop conf keys set on `hadoopConf` for which [[isGcsAuthKey]] holds, full key + * names only -- NEVER their values, which are credential material and must never enter a + * decline reason. Iterates the conf map directly: no provider resolution, no I/O. + */ + private def gcsAuthKeys(hadoopConf: Configuration): Seq[String] = + hadoopConf + .iterator() + .asScala + .collect { + case entry + if isGcsAuthKey(entry.getKey) && entry.getValue != null && + entry.getValue.nonEmpty => + entry.getKey + } + .toSeq + .distinct + .sorted + + /** + * Decline reason when any of `uris` resolves to a `gs://` authority AND `hadoopConf` sets any + * key [[isGcsAuthKey]] flags, or `None` when claimable. Native forwards none of `fs.gs.*` (nor + * any of the legacy/deprecated `google.cloud.*` namespaces) to the object store, so a scan + * relying solely on Hadoop-side GCS credentials would claim here but then fail authentication + * natively. Application Default Credentials work identically in both engines and need no Hadoop + * conf key, so an ADC-only configuration still claims. Never interpolates a resolved value, + * only key names. + */ + private[delta] def gcsHadoopOnlyAuthReason( + hadoopConf: Configuration, + uris: Seq[URI]): Option[String] = { + val gcsUris = uris.filter(isGcsScheme) + if (gcsUris.isEmpty) { + return None + } + val authKeys = gcsAuthKeys(hadoopConf) + if (authKeys.isEmpty) { + return None + } + val buckets = gcsUris.flatMap(gcsBucket).distinct.sorted + Some( + "Native Delta scan does not support GCS authentication configured only via Hadoop conf " + + s"key(s) ${authKeys.mkString(", ")} for gs://${buckets.mkString(", gs://")} " + + "(the native GCS client does not forward fs.gs.* options; only Application Default " + + "Credentials -- environment or metadata-server -- are available natively)") + } + + /** + * True when `dataType` is, or structurally contains (through array elements or map keys/ + * values), a [[StructType]]. Only [[StructType]] fields carry Delta's physical, column-mapped + * names; array/map labels themselves are never column-mapped. + */ + private def containsNestedStruct(dataType: DataType): Boolean = dataType match { + case _: StructType => true + case ArrayType(elementType, _) => containsNestedStruct(elementType) + case MapType(keyType, valueType, _) => + containsNestedStruct(keyType) || containsNestedStruct(valueType) + case _ => false + } + + /** + * True when `node` is a positional-output union -- `UnionExec` or `CometUnionExec`. Both + * compute output positionally from the FIRST child's attributes, so a value carried only by a + * LATER branch needs an explicit positional walk below. Compared by class name (the + * [[isDeltaScan]] idiom) to avoid a compile-time dependency; an unmatched name is still safe, + * caught by the generic child-output safety net below. + */ + private def isPositionalUnion(node: SparkPlan): Boolean = { + val name = node.getClass.getSimpleName + name == "UnionExec" || name == "CometUnionExec" + } + + /** + * True when the scan's row-index column value is provably dead above the scan. The standard DV + * plan shape routes it only into a `named_struct(... row_index ...) AS _metadata` projection + * whose result the final projection discards; anything else (a query actually selecting + * `_metadata.row_index`, OR a write sink -- `DataWritingCommandExec`, `WriteFilesExec`, a DSv2 + * `V2TableWriteExec` -- persisting it) makes the value live and must decline. Conservative: any + * unrecognized consumption pattern returns false. + */ + private def rowIndexUnusedAbove(plan: SparkPlan, scanExec: FileSourceScanExec): Boolean = + attributesUnusedAbove( + plan, + scanExec, + scanExec.output.filter(_.name == CometDeltaNativeScan.RowIndexColumn).map(_.exprId).toSet) + + /** + * True when the values of the scan output attributes `attrIds` are provably dead above the + * scan, by the analysis described on [[rowIndexUnusedAbove]]. + */ + private def attributesUnusedAbove( + plan: SparkPlan, + scanExec: FileSourceScanExec, + attrIds: Set[ExprId]): Boolean = { + if (attrIds.isEmpty) { + return true + } + // Transitive taint analysis: everything derived from the seed attributes within the + // visible plan, via Project aliases or positionally across a union. The plan may be an AQE + // stage fragment, so tainted values escaping to the fragment's own output must decline too. + var tainted = attrIds + var changed = true + while (changed) { + changed = false + plan.foreach { + case p: ProjectExec => + p.projectList.foreach { + case a: Alias + if !tainted.contains(a.exprId) && + a.references.exists(r => tainted.contains(r.exprId)) => + tainted += a.exprId + changed = true + case _ => + } + case u if isPositionalUnion(u) => + // Output attributes carry the FIRST child's expression IDs, so a value tainted only in + // a LATER branch is otherwise invisible; walk it forward positionally instead. + // `children` can be re-parented by AQE after `output` is frozen, so an arity mismatch on + // ANY child (which would make a positional zip silently truncate) forces a decline. + if (u.children.exists(_.output.length != u.output.length)) { + return false + } + u.children.foreach { child => + child.output.zip(u.output).foreach { + case (from, to) if tainted.contains(from.exprId) && !tainted.contains(to.exprId) => + tainted += to.exprId + changed = true + case _ => + } + } + case _ => + } + } + val nonProjectConsumer = plan.exists { + case _: ProjectExec => false + case n if n ne scanExec => + n.expressions.exists(_.references.exists(r => tainted.contains(r.exprId))) + case _ => false + } + val escapes = plan.output.exists(a => tainted.contains(a.exprId)) + // Generic safety net for every OTHER node, of ANY arity (joins and other multi-child + // shapes, but also plain one-child nodes; positional unions and Project are exempt, already + // handled precisely above -- Project's own output legitimately omits a tainted attribute it + // dropped, which is not a leak). A tainted attribute a child contributes must either survive + // into the node's own output under the SAME expression ID or be consumed by one of the + // node's own expressions; otherwise decline. This catches two shapes: a multi-child node + // dropping the side carrying the tainted attribute (e.g. a LEFT SEMI/ANTI join), and a + // one-child WRITE SINK -- DataWritingCommandExec, WriteFilesExec, and the DSv2 + // AppendDataExec/OverwriteByExpressionExec/... family (V2TableWriteExec) -- that executes + // its child purely for the side effect of persisting its rows and so has an EMPTY output of + // its own. Such a sink neither preserves the tainted attribute (nothing survives into an + // empty output) nor references it in an expression, so without this check it looks like an + // inert pass-through even though the write persists whatever value the reader returned, + // including a DV scan's dead synthetic row-index constant. + val childOutputLeak = plan.exists { + case u if isPositionalUnion(u) => false + case _: ProjectExec => false + case n if n.children.nonEmpty => + n.children.exists { c => + c.output.exists { attr => + tainted.contains(attr.exprId) && + !n.output.exists(_.exprId == attr.exprId) && + !n.expressions.exists(_.references.exists(_.exprId == attr.exprId)) + } + } + case _ => false + } + !nonProjectConsumer && !escapes && !childOutputLeak + } +} diff --git a/contrib/delta-spark/src/main/scala/org/apache/comet/contrib/delta/DeltaSparkConfigProvider.scala b/contrib/delta-spark/src/main/scala/org/apache/comet/contrib/delta/DeltaSparkConfigProvider.scala new file mode 100644 index 00000000000..c0a4f5e0fdb --- /dev/null +++ b/contrib/delta-spark/src/main/scala/org/apache/comet/contrib/delta/DeltaSparkConfigProvider.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.contrib.delta + +import org.apache.comet.{CometConfigProvider, ConfigEntry} + +/** + * Exposes this contrib's config entries to `GenerateDocs`. Note: with the current module layout + * (`contrib/delta-spark` depends on `comet-spark`) the doc build cannot see this provider; it + * exists to satisfy the contrib-conf contract and becomes active if the module is ever folded + * into the spark build like `contrib/delta` is. + */ +class DeltaSparkConfigProvider extends CometConfigProvider { + override def configs: Seq[ConfigEntry[_]] = DeltaScanConf.all + override def docPage: String = "delta.md" + override def docCategory: String = "delta" +} diff --git a/contrib/delta-spark/src/main/scala/org/apache/comet/contrib/delta/DeltaSparkScanEnvelope.scala b/contrib/delta-spark/src/main/scala/org/apache/comet/contrib/delta/DeltaSparkScanEnvelope.scala new file mode 100644 index 00000000000..5bfad727b2e --- /dev/null +++ b/contrib/delta-spark/src/main/scala/org/apache/comet/contrib/delta/DeltaSparkScanEnvelope.scala @@ -0,0 +1,54 @@ +/* + * 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.contrib.delta + +import org.apache.comet.serde.OperatorOuterClass +import org.apache.comet.serde.OperatorOuterClass.Operator + +/** + * Packs and unpacks the JVM-planned `DeltaSparkScan` message in core's generic `ContribScan` + * envelope (`contrib_scan` on `Operator`). The native dispatcher routes by `type_url`, so this + * contrib's identifier is the only coupling between the JVM and native sides; core names no Delta + * type. + */ +object DeltaSparkScanEnvelope { + + /** + * Contrib-owned identifier for the message, mirrored by `DELTA_SPARK_SCAN_TYPE_NAME` in + * native's `delta_spark_scan.rs`. Distinct from the kernel path's + * `comet.contrib.delta.DeltaScan`. + */ + val TypeUrl = "type.googleapis.com/comet.contrib.delta_spark.DeltaSparkScan" + + def pack(scan: OperatorOuterClass.DeltaSparkScan): OperatorOuterClass.ContribScan = + OperatorOuterClass.ContribScan + .newBuilder() + .setTypeUrl(TypeUrl) + .setValue(scan.toByteString) + .build() + + /** Whether this operator carries this contrib's scan (and not some other contrib's). */ + def matches(op: Operator): Boolean = + op.hasContribScan && op.getContribScan.getTypeUrl == TypeUrl + + /** Callers must check `matches` first. */ + def unpack(op: Operator): OperatorOuterClass.DeltaSparkScan = + OperatorOuterClass.DeltaSparkScan.parseFrom(op.getContribScan.getValue) +} diff --git a/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/CometDeltaNativeScanExec.scala b/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/CometDeltaNativeScanExec.scala new file mode 100644 index 00000000000..c1737172dd2 --- /dev/null +++ b/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/CometDeltaNativeScanExec.scala @@ -0,0 +1,329 @@ +/* + * 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.rdd.RDD +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.plans.QueryPlan +import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, UnknownPartitioning} +import org.apache.spark.sql.execution.{FileSourceScanExec, InSubqueryExec, ReusedSubqueryExec, ScalarSubquery, SparkPlan, SubqueryAdaptiveBroadcastExec} +import org.apache.spark.sql.execution.datasources.HadoopFsRelation +import org.apache.spark.sql.execution.metric.SQLMetric +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.vectorized.ColumnarBatch + +import org.apache.comet.contrib.delta.DeltaSparkScanEnvelope +import org.apache.comet.serde.OperatorOuterClass +import org.apache.comet.serde.OperatorOuterClass.Operator + +/** + * Native scan node for Delta Lake tables (contrib). Delta's own planning (log replay, snapshot + * resolution, partition pruning) has already run inside delta-spark by the time this node is + * created from the DSv1 [[FileSourceScanExec]]; file listing and split planning are delegated to + * a [[CometScanExec]] helper, and data reads execute through Comet's native DataFusion parquet + * machinery, inheriting row-group and page-index pruning. + * + * DPP: `runtimeFilters` is a constructor field included in equality, so its rewrite (via + * [[CometScanWithPlanData]]) survives plan copies -- a transient field would be dropped by + * `TreeNode.makeCopy` on MERGE re-planning (the CometIcebergNativeScanExec lesson). + */ +case class CometDeltaNativeScanExec( + override val nativeOp: Operator, + override val output: Seq[Attribute], + requiredSchema: StructType, + runtimeFilters: Seq[Expression], + dataFilters: Seq[Expression], + @transient relation: HadoopFsRelation, + originalPlan: FileSourceScanExec, + override val serializedPlanOpt: SerializedPlan, + sourceKey: String) + extends CometLeafExec + with CometScanWithPlanData { + + override val nodeName: String = s"CometDeltaNativeScan $relation" + + override def sparkFallback(newChildren: Seq[SparkPlan]): SparkPlan = { + // AQE rewrites DPP placeholders in runtimeFilters, not in originalPlan, so restore the scan + // with the live partition filters. dataFilters stays as Spark planned it: the live list also + // carries subquery filters harvested from the parent FilterExec. + val restoredScan = originalPlan.copy(partitionFilters = runtimeFilters) + originalPlan.logicalLink.foreach(restoredScan.setLogicalLink) + restoredScan + } + + // Derived from (originalPlan, runtimeFilters), never stored: any copy of this node + // automatically gets a helper consistent with ITS runtimeFilters, avoiding the #3510 class of + // bug where a stored helper field desyncs from rewritten filters. Costs one extra file listing + // per executed instance; correctness over the duplicate driver-side listing. + // + // Forcing invariant: this lazy val is forced by the `metrics` override below, and AQE's UI + // plan-walk calls `.metrics` on every node MID-PLANNING, including while a DPP subquery is + // still an adaptive placeholder or a partition filter holds an unresolved ScalarSubquery (see + // `hasUnevaluableSubqueryFilter` below). That's safe ONLY because constructing `scanHelper` is a + // cheap case-class build with no file listing, and core's `CometScanExec.metrics` touches only + // `wrapped.driverMetrics` (populated by Spark's own planning) plus a static metric-node + // constructor -- neither file listing nor subquery resolution. If core's `metrics` ever touches + // either, forcing `scanHelper` here would resurrect the AQE mid-planning crashes this invariant + // prevents. + @transient private lazy val scanHelper: CometScanExec = + CometDeltaNativeScanExec.planningHelper(originalPlan, runtimeFilters) + + // NOT lazy val: while a DPP subquery is still an adaptive placeholder, or a partition filter + // holds an unresolved scalar subquery, this returns a temporary value that must not be + // memoized -- after CometPlanAdaptiveDynamicPruningFilters rewrites the filters (DPP case) or + // AQE resolves the subquery (scalar case), later reads must see the real post-pruning + // partition count. + override def outputPartitioning: Partitioning = + if (hasUnevaluableSubqueryFilter) UnknownPartitioning(0) + else UnknownPartitioning(perPartitionData.length) + + // runtimeFilters IS scanHelper.partitionFilters element-for-element, so checking runtimeFilters + // here avoids constructing/forcing the derived scanHelper just to read partitioning. The + // InSubqueryExec placeholder shapes mirror + // CometPlanAdaptiveDynamicPruningFilters.extractSABData + hasWrappedSAB -- keep in sync. The + // ScalarSubquery case is probed rather than treated as permanently unevaluable: Spark exposes no + // public finished/updated flag on ExecSubqueryExpression, but `eval()` doubles as one -- it only + // reads the cached `result` behind a `require(updated, ...)` guard, while the subquery is + // actually run by `updateResult()` (invoked separately during prepare/AQE), never by `eval()`. + // Once resolved, outputPartitioning below reports the real perPartitionData.length instead of + // staying at zero -- a fused native parent's buildNativeContext requires that count to match. + private def hasUnevaluableSubqueryFilter: Boolean = + runtimeFilters.exists(_.exists { + // Match `e: InSubqueryExec` and dispatch on e.plan rather than unapplying InSubqueryExec + // directly: its unapply arity differs across Spark versions and this module ships no + // version shim. + case e: InSubqueryExec => isAdaptivePlaceholder(e.plan) + case s: ScalarSubquery => !isScalarSubqueryResolved(s) + case _ => false + }) + + // `eval()` never triggers the subquery's execution: on a resolved subquery it is a pure cached + // read of `result` (verified against bytecode: `Predef.require(updated(), ...)` then a plain + // field read), so this probe is safe to call repeatedly, including from AQE's mid-planning plan + // walks. Pre-resolution, the ONLY throw is `require`'s `IllegalArgumentException`; catch exactly + // that, since anything else escaping is a genuine bug we must not mask as unpartitioned. + private def isScalarSubqueryResolved(s: ScalarSubquery): Boolean = + try { + s.eval() + true + } catch { + case _: IllegalArgumentException => false + } + + private def isAdaptivePlaceholder(p: SparkPlan): Boolean = p match { + case ReusedSubqueryExec(inner) => isAdaptivePlaceholder(inner) + case _: CometSubqueryAdaptiveBroadcastExec => true + case _: SubqueryAdaptiveBroadcastExec => true + case _ => false + } + + override lazy val outputOrdering: Seq[SortOrder] = originalPlan.outputOrdering + + override def dynamicPruningFilters: Seq[Expression] = runtimeFilters + + override def withDynamicPruningFilters(filters: Seq[Expression]): SparkPlan = { + // A real copy: runtimeFilters is a constructor field included in equality, so the copy + // survives enclosing-block rebuilds, and the derived scanHelper picks up the rewritten + // filters automatically. + copy(runtimeFilters = filters) + } + + /** + * Lazy split-mode serialization, mirroring CometNativeScanExec: common data was serialized at + * planning; per-partition file lists serialize here, at execution time. + */ + @transient private lazy val serializedPartitionData + : (Array[Byte], Array[Array[Byte]], Array[Seq[String]]) = { + // Resolve the helper's DPP subqueries: it holds its own InSubqueryExec instances that + // Spark's expressions walk does not see (the helper is derived, not a child). + scanHelper.partitionFilters.foreach { + case DynamicPruningExpression(e: InSubqueryExec) if e.values().isEmpty => + e.updateResult() + case _ => + } + + val commonBytes = { + val deltaScan = DeltaSparkScanEnvelope.unpack(nativeOp) + // Scalar subqueries in dataFilters were unresolved at planning; resolve them now and + // append them as pushed filters, as CometNativeScanExec.serializedPartitionData does. + // has_data_filters follows their presence, not the serialized count: a filter that fails + // to serialize still keeps native on the safe timestamp conversion for a filtered scan. + val resolved = org.apache.comet.contrib.delta.CometDeltaNativeScan + .resolvedSubqueryFilters(dataFilters, output, requiredSchema, conf) + val common = if (!resolved.hasResolvedFilters) { + deltaScan.getCommon + } else { + val builder = deltaScan.getCommon.toBuilder + builder.setHasDataFilters(true) + resolved.protos.foreach(builder.addDataFilters) + builder.build() + } + OperatorOuterClass.DeltaSparkScan + .newBuilder() + .setCommon(common) + .setDeltaCommon(deltaScan.getDeltaCommon) + .build() + .toByteArray + } + + val filePartitions = scanHelper.getFilePartitions() + + val tableRoot = DeltaSparkScanEnvelope.unpack(nativeOp).getDeltaCommon.getTableRoot + val perPartitionBytes = filePartitions.map { filePartition => + org.apache.comet.contrib.delta.CometDeltaNativeScan + .serializePartition(filePartition, originalPlan, tableRoot) + }.toArray + + val perPartitionPaths = filePartitions.map(_.files.map(_.filePath.toString).toSeq).toArray + + (commonBytes, perPartitionBytes, perPartitionPaths) + } + + override def commonData: Array[Byte] = serializedPartitionData._1 + + override def perPartitionData: Array[Array[Byte]] = serializedPartitionData._2 + + def perPartitionFilePaths: Array[Seq[String]] = serializedPartitionData._3 + + override def doExecuteColumnar(): RDD[ColumnarBatch] = { + val nativeMetrics = CometMetricNode.fromCometPlan(this) + val serializedPlan = CometExec.serializeNativePlan(nativeOp) + + new CometExecRDD( + sparkContext, + Seq.empty, + Map(sourceKey -> commonData), + Map(sourceKey -> perPartitionData), + serializedPlan, + PlanDataInjector.planFingerprint(serializedPlan), + perPartitionData.length, + output.length, + nativeMetrics, + Seq.empty, + None, + Seq.empty, + perPartitionFilePaths = perPartitionFilePaths, + outputTypes = output.map(_.dataType), + reportScanInputMetrics = true) + } + + override def doCanonicalize(): CometDeltaNativeScanExec = { + val canonOriginal = if (originalPlan != null) { + val stripped = originalPlan.copy(partitionFilters = + CometScanUtils.filterUnusedDynamicPruningExpressions(originalPlan.partitionFilters)) + stripped.doCanonicalize() + } else { + null + } + CometDeltaNativeScanExec( + nativeOp, + output.map(QueryPlan.normalizeExpressions(_, output)), + requiredSchema, + QueryPlan.normalizePredicates( + CometScanUtils.filterUnusedDynamicPruningExpressions(runtimeFilters), + output), + QueryPlan.normalizePredicates(dataFilters, output), + relation, + canonOriginal, + SerializedPlan(None), + "") + } + + override def stringArgs: Iterator[Any] = Iterator(output, runtimeFilters) + + override def equals(obj: Any): Boolean = obj match { + case other: CometDeltaNativeScanExec => + this.originalPlan == other.originalPlan && + this.serializedPlanOpt == other.serializedPlanOpt && + this.runtimeFilters == other.runtimeFilters && + this.dataFilters == other.dataFilters + case _ => false + } + + override def hashCode(): Int = + java.util.Objects.hash(originalPlan, serializedPlanOpt, runtimeFilters, dataFilters) + + private val driverMetricKeys = + Set( + "numFiles", + "filesSize", + "numPartitions", + "metadataTime", + "staticFilesNum", + "staticFilesSize", + "pruningTime") + + // Forces `scanHelper` (see its doc above for why that -- and reading `.metrics` off it -- is + // safe even when AQE calls `.metrics` mid-planning against an unresolved DPP/scalar subquery). + override lazy val metrics: Map[String, SQLMetric] = { + CometMetricNode.nativeScanMetrics(session.sparkContext) ++ + scanHelper.metrics.filter { case (k, _) => driverMetricKeys.contains(k) } + } +} + +object CometDeltaNativeScanExec { + + /** + * File-planning helper: reuses CometScanExec's listing/splitting/DPP machinery. Files with a + * deletion vector are split like any other file: a claimed scan requires Delta's reader + * optimizations to be enabled, which is exactly what DeltaParquetFileFormat.isSplitable + * returns, and Spark's row-index split gate does not apply to a claimed scan. Each split then + * fetches and decodes the whole deletion vector, reads the footer and admits memory against a + * bound for the whole file, but builds row selectors only for the row groups that start inside + * the split, which are the row groups the reader keeps, and shrinks the reservation it holds to + * those selectors plus the whole-file page-index allowance. + */ + def planningHelper( + scanExec: FileSourceScanExec, + partitionFilters: Seq[Expression]): CometScanExec = + CometScanExec( + scanExec.relation, + scanExec.output, + scanExec.requiredSchema, + partitionFilters, + scanExec.optionalBucketSet, + scanExec.optionalNumCoalescedBuckets, + scanExec.dataFilters, + scanExec.tableIdentifier, + scanExec.disableBucketedScan, + scanExec) + + def apply( + nativeOp: Operator, + scanExec: FileSourceScanExec, + subqueryDataFilters: Seq[Expression] = Seq.empty): CometDeltaNativeScanExec = { + // subqueryDataFilters: subquery predicates harvested from the covering FilterExec at claim + // time (Spark 3.x keeps them out of scanExec.dataFilters; see + // CometDeltaNativeScan.subqueryFiltersFromParent). Carried in dataFilters so the + // execution-time resolve-and-push path sees them; correctness never depends on them. + val exec = CometDeltaNativeScanExec( + nativeOp, + scanExec.output, + scanExec.requiredSchema, + scanExec.partitionFilters, + scanExec.dataFilters ++ subqueryDataFilters, + scanExec.relation, + scanExec, + SerializedPlan(None), + DeltaSparkScanEnvelope.unpack(nativeOp).getDeltaCommon.getSourceKey) + scanExec.logicalLink.foreach(exec.setLogicalLink) + exec + } +} diff --git a/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/DeltaPlanDataInjector.scala b/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/DeltaPlanDataInjector.scala new file mode 100644 index 00000000000..6697b57bde8 --- /dev/null +++ b/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/DeltaPlanDataInjector.scala @@ -0,0 +1,89 @@ +/* + * 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 scala.jdk.CollectionConverters._ + +import org.apache.comet.contrib.delta.DeltaSparkScanEnvelope +import org.apache.comet.serde.{OperatorOuterClass, QueryContextInterner} +import org.apache.comet.serde.OperatorOuterClass.Operator + +/** + * PlanDataInjector for the Delta contrib scan, discovered by core's ServiceLoader (see the + * `META-INF/services` resource). Lives in this package because [[PlanDataInjector]] is + * `private[comet]`. + */ +class DeltaPlanDataInjector extends PlanDataInjector { + + // The partition-invariant half of a DeltaSparkScan: common + delta_common, no file partition. + override type Prepared = OperatorOuterClass.DeltaSparkScan + + override val opStructCase: Operator.OpStructCase = Operator.OpStructCase.CONTRIB_SCAN + + override def canInject(op: Operator): Boolean = + DeltaSparkScanEnvelope.matches(op) && { + val scan = DeltaSparkScanEnvelope.unpack(op) + scan.hasCommon && !scan.hasFilePartition + } + + override def getKey(op: Operator): Option[String] = + Some(DeltaSparkScanEnvelope.unpack(op).getDeltaCommon.getSourceKey) + + // commonBytes is a DeltaSparkScan proto carrying common + delta_common (no file partition). + // Parsing it dominates inject() on wide schemas; injectPlanData memoizes the result per stage. + override def prepareCommon(commonBytes: Array[Byte]): Prepared = + OperatorOuterClass.DeltaSparkScan.parseFrom(commonBytes) + + override def inject(op: Operator, common: Prepared, partitionBytes: Array[Byte]): Operator = { + // partitionBytes is a DeltaSparkScan proto carrying only this partition's file list. + val partitionOnly = OperatorOuterClass.DeltaSparkScan.parseFrom(partitionBytes) + + val scanBuilder = OperatorOuterClass.DeltaSparkScan + .newBuilder() + .setCommon(common.getCommon) + .setDeltaCommon(common.getDeltaCommon) + .setFilePartition(partitionOnly.getFilePartition) + + op.toBuilder.setContribScan(DeltaSparkScanEnvelope.pack(scanBuilder.build())).build() + } +} + +object DeltaPlanDataInjector { + + /** + * The key under which a Delta scan's planning data is stored and looked up. Written into + * `DeltaSparkScanCommon.source_key` on the driver and read back by + * [[DeltaPlanDataInjector.getKey]] on the executor, so both sides agree by construction. + * Mirrors `NativeScanPlanDataInjector.sourceKey` (source string carries the plan node id, so + * two scans of the same table in one plan, self-join, MERGE, get distinct keys), plus the table + * root for extra safety across tables with identical projections. + */ + def sourceKey(tableRoot: String, common: OperatorOuterClass.NativeScanCommon): String = { + val dataFilters = common.getDataFiltersList.asScala + .map(QueryContextInterner.stripQueryContexts(_).toString) + val keyComponents = Seq( + tableRoot, + common.getRequiredSchemaList.toString, + dataFilters.mkString("[", ", ", "]"), + common.getProjectionVectorList.toString, + common.getFieldsList.toString) + s"delta_${common.getSource}_${keyComponents.mkString("|").hashCode}" + } +} diff --git a/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaDmlReproSuite.scala b/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaDmlReproSuite.scala new file mode 100644 index 00000000000..8aa1c0643eb --- /dev/null +++ b/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaDmlReproSuite.scala @@ -0,0 +1,155 @@ +/* + * 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.contrib.delta + +import scala.collection.mutable.ListBuffer + +import org.apache.spark.CometListenerBusUtils +import org.apache.spark.sql.delta.DeltaLog +import org.apache.spark.sql.execution.{FileSourceScanExec, QueryExecution, SparkPlan} +import org.apache.spark.sql.util.QueryExecutionListener + +import org.apache.comet.ExtendedExplainInfo + +/** + * Repro for Delta's own DeletionVectorsSuite expectation: DELETE on a DV-enabled table must WRITE + * deletion vectors (not rewrite files) with Comet active. Mirrors "DELETE with DVs - on a table + * with no prior DVs". + */ +class CometDeltaDmlReproSuite extends CometDeltaTestBase { + + /** + * Every [[SparkPlan]] Delta's own internal DataFrame actions executed during `body`, captured + * via a [[QueryExecutionListener]] rather than the outer statement's own plan: Delta's DML + * commands (DELETE/UPDATE/MERGE) drive `findTouchedFiles` through separate internal + * `collect`/`count` actions on their own [[QueryExecution]]s, invisible to `df.queryExecution` + * on the outer SQL statement. + */ + private def capturePlansDuring(body: => Unit): Seq[SparkPlan] = { + val plans = ListBuffer.empty[SparkPlan] + val listener = new QueryExecutionListener { + override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = { + plans += qe.executedPlan + } + override def onFailure( + funcName: String, + qe: QueryExecution, + exception: Exception): Unit = {} + } + spark.listenerManager.register(listener) + try { + body + // The listener bus delivers asynchronously, so the plans are not all in hand until it has + // drained. + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + } finally { + spark.listenerManager.unregister(listener) + } + plans.toSeq + } + + test( + "DELETE's internal deletion-vector-generating scan declines the row-index-outside-a-DV-" + + "scan reason (the read-side counterpart of the DV-write repro above)") { + withSQLConf("spark.databricks.delta.properties.defaults.enableDeletionVectors" -> "true") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 1000, 1, 4).write.format("delta").save(path) + + val capturedPlans = capturePlansDuring { + spark.sql(s"DELETE FROM delta.`$path` WHERE id % 2 = 0 AND id < 200") + } + + // Before writing a deletion vector, DELETE must first learn WHICH rows matched the + // predicate, so it reads each candidate file's `_metadata.row_index` directly (a bare + // row-index column, with no `is_row_deleted` alongside it -- unlike a normal DV-applying + // read, no existing DV is applied to this scan, since the very DV being computed does not + // exist yet). DeltaScanSupport.declineReason's hasRowIndex-without-hasIsRowDeleted gate + // exists precisely to keep this bookkeeping scan on Spark's reader: claiming it with a + // dead constant row-index would feed wrong (constant) row indexes into the DV this DELETE + // is trying to build. This must remain a plain Spark FileSourceScanExec here, never a + // CometDeltaNativeScanExec. + val declinedRowIndexScans = capturedPlans.flatMap { plan => + collectWithSubqueries(stripAQEPlan(plan)) { + case f: FileSourceScanExec + if DeltaScanSupport.isDeltaScan(f) && + f.requiredSchema.exists(_.name == CometDeltaNativeScan.RowIndexColumn) && + !f.requiredSchema.exists(_.name == CometDeltaNativeScan.IsRowDeletedColumn) => + f + } + } + assert( + declinedRowIndexScans.nonEmpty, + "expected to observe at least one internal row-index-only scan while DELETE " + + "computed which rows to mark in the new deletion vector") + + val reasons = + declinedRowIndexScans.flatMap(f => new ExtendedExplainInfo().getFallbackReasons(f)) + assert( + reasons.exists(_.contains("row-index reads outside a deletion-vector scan")), + "expected the internal row-index scan to carry the row-index-outside-a-DV-scan " + + s"decline reason, got: ${reasons.mkString(", ")}") + + val log = DeltaLog.forTable(spark, path) + val withDvs = log.update().allFiles.collect().count(_.deletionVector != null) + assert(withDvs > 0, s"expected at least one file to have a DV written, got $withDvs") + assert(spark.read.format("delta").load(path).count() == 900) + } + } + } + + test("DELETE writes DVs with useMetadataRowIndex=true (metadata row-index DML shape)") { + withSQLConf( + "spark.databricks.delta.properties.defaults.enableDeletionVectors" -> "true", + "spark.databricks.delta.deletionVectors.useMetadataRowIndex" -> "true") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 1000, 1, 500).write.format("delta").save(path) + spark.sql(s"DELETE FROM delta.`$path` WHERE id % 2 = 0 AND id < 200") + + val log = DeltaLog.forTable(spark, path) + val withDvs = log.update().allFiles.collect().count(_.deletionVector != null) + assert(withDvs == 100, s"expected 100 files with DVs, got $withDvs") + assert(spark.read.format("delta").load(path).count() == 900) + } + } + } + + test("DELETE writes DVs rather than rewriting files") { + withSQLConf( + "spark.databricks.delta.properties.defaults.enableDeletionVectors" -> "true", + "spark.databricks.delta.delete.deletionVectors.persistent" -> "true") { + withTempDir { base => + // Mirror Delta's DeletionVectorsTestUtils: paths with spaces and a literal %2a. + val dir = new java.io.File(base, "s p a r k %2a") + val path = dir.getAbsolutePath + spark.range(0, 1000, 1, 500).write.format("delta").save(path) + spark.sql(s"DELETE FROM delta.`$path` WHERE id % 2 = 0 AND id < 200") + + val log = DeltaLog.forTable(spark, path) + val files = log.update().allFiles.collect() + val withDvs = files.count(_.deletionVector != null) + assert(files.length == 500, s"expected 500 files, got ${files.length}") + assert(withDvs == 100, s"expected 100 files with DVs, got $withDvs") + assert(spark.read.format("delta").load(path).count() == 900) + } + } + } +} diff --git a/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala b/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala new file mode 100644 index 00000000000..3b82cb32654 --- /dev/null +++ b/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala @@ -0,0 +1,3952 @@ +/* + * 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.contrib.delta + +import java.io.File + +import scala.collection.mutable +import scala.collection.mutable.ListBuffer + +import org.apache.spark.CometListenerBusUtils +import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} +import org.apache.spark.sql.{DataFrame, Row} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, DynamicPruningExpression, Expression, NamedExpression, StructsToJson} +import org.apache.spark.sql.comet.{CometDeltaNativeScanExec, CometSubqueryBroadcastExec} +import org.apache.spark.sql.execution.{FileSourceScanExec, InSubqueryExec, QueryExecution, ReusedSubqueryExec, ScalarSubquery, SparkPlan, SubqueryBroadcastExec, SubqueryExec} +import org.apache.spark.sql.execution.datasources.v2.V2TableWriteExec +import org.apache.spark.sql.functions.{col, lit, to_json} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{ByteType, LongType, StringType, StructField, StructType} +import org.apache.spark.sql.util.QueryExecutionListener + +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus} +import org.apache.comet.ExtendedExplainInfo +import org.apache.comet.serde.OperatorOuterClass +import org.apache.comet.serde.operator.CometNativeScan + +/** + * Differential suite: append-only Delta tables read through the native Delta scan must produce + * results identical to Spark's Delta reader, engage the native operator, and prune at row-group + * and page level. + */ +class CometDeltaNativeScanSuite extends CometDeltaTestBase { + + test("plain delta table reads natively with identical results") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark + .range(0, 1000) + .selectExpr("id", "id * 2 as v", "cast(id as string) as s") + .write + .format("delta") + .save(path) + + val df = spark.read.format("delta").load(path).filter(col("id") > 500) + checkDeltaNativeScanAnswer(df) + } + } + + test("projection and filter on delta table") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark + .range(0, 1000) + .selectExpr("id", "id % 10 as bucket", "cast(id as double) as d") + .write + .format("delta") + .save(path) + + val df = spark.read + .format("delta") + .load(path) + .select("bucket", "d") + .filter(col("d") < 100.0) + checkDeltaNativeScanAnswer(df) + } + } + + test("partitioned delta table with partition filter") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark + .range(0, 1000) + .selectExpr("id", "id % 7 as p") + .write + .format("delta") + .partitionBy("p") + .save(path) + + val df = spark.read.format("delta").load(path).filter(col("p") === 3) + checkDeltaNativeScanAnswer(df) + assert(df.count() > 0) + } + } + + test("multi-file delta table after several appends") { + withTempPath { dir => + val path = dir.getAbsolutePath + for (i <- 0 until 4) { + spark + .range(i * 100L, (i + 1) * 100L) + .selectExpr("id", "id * 3 as v") + .write + .format("delta") + .mode("append") + .save(path) + } + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + assert(df.count() == 400) + } + } + + test("time travel VERSION AS OF reads natively") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 100).write.format("delta").save(path) + spark.range(100, 200).write.format("delta").mode("append").save(path) + + val v0 = spark.read.format("delta").option("versionAsOf", 0).load(path) + checkDeltaNativeScanAnswer(v0) + assert(v0.count() == 100) + } + } + + test("selective predicate prunes row groups and pages") { + withTempPath { dir => + val path = dir.getAbsolutePath + // Small row groups + page-level stats: sorted data so min/max stats are tight. The Delta + // writer ignores parquet.* DataFrameWriter options, so set them on the Hadoop conf. + val hadoopConf = spark.sparkContext.hadoopConfiguration + val oldBlockSize = hadoopConf.get("parquet.block.size") + val oldPageSize = hadoopConf.get("parquet.page.size") + hadoopConf.setInt("parquet.block.size", 256 * 1024) + hadoopConf.setInt("parquet.page.size", 16 * 1024) + try { + spark + .range(0, 500000) + .selectExpr("id", "id * 2 as v") + .sort("id") + .coalesce(1) + .write + .format("delta") + .save(path) + } finally { + if (oldBlockSize == null) hadoopConf.unset("parquet.block.size") + else hadoopConf.set("parquet.block.size", oldBlockSize) + if (oldPageSize == null) hadoopConf.unset("parquet.page.size") + else hadoopConf.set("parquet.page.size", oldPageSize) + } + + def query = spark.read + .format("delta") + .load(path) + .filter(col("id") >= 100 && col("id") < 200) + checkDeltaNativeScanAnswer(query) + + // checkSparkAnswer re-plans the query, so read metrics from a DataFrame we execute + // ourselves (collect() runs THIS Dataset's queryExecution; count() would plan a new one): + // its executed plan holds the metric objects native execution updated. + val df = query + assert(df.collect().length == 100) + val scans = deltaNativeScans(df) + assert(scans.size == 1) + val metrics = scans.head.metrics + val rowGroupsPruned = metrics.get("row_groups_pruned_statistics").map(_.value).getOrElse(0L) + val pagesPruned = metrics.get("page_index_rows_pruned").map(_.value).getOrElse(0L) + assert( + rowGroupsPruned > 0, + s"expected row-group pruning; metrics: ${metrics.map { case (k, v) => s"$k=${v.value}" }}") + assert( + pagesPruned > 0, + s"expected page-index pruning; metrics: ${metrics.map { case (k, v) => + s"$k=${v.value}" + }}") + } + } + + test("scalar subquery data filter is pushed down and prunes row groups and pages") { + withTempPath { dir => + val path = s"${dir.getAbsolutePath}/data" + val thresholds = s"${dir.getAbsolutePath}/thresholds" + // Same layout as the selective-predicate test: small row groups + tight page stats. + val hadoopConf = spark.sparkContext.hadoopConfiguration + val oldBlockSize = hadoopConf.get("parquet.block.size") + val oldPageSize = hadoopConf.get("parquet.page.size") + hadoopConf.setInt("parquet.block.size", 256 * 1024) + hadoopConf.setInt("parquet.page.size", 16 * 1024) + try { + spark + .range(0, 500000) + .selectExpr("id", "id * 2 as v") + .sort("id") + .coalesce(1) + .write + .format("delta") + .save(path) + } finally { + if (oldBlockSize == null) hadoopConf.unset("parquet.block.size") + else hadoopConf.set("parquet.block.size", oldBlockSize) + if (oldPageSize == null) hadoopConf.unset("parquet.page.size") + else hadoopConf.set("parquet.page.size", oldPageSize) + } + spark + .sql("SELECT CAST(100 AS BIGINT) AS lo, CAST(200 AS BIGINT) AS hi") + .write + .format("delta") + .save(thresholds) + + // Scalar subqueries are PlanExpressions: unresolved at planning, so the bounds can + // only reach the native reader via the execution-time resolve-and-append path. + def query = spark.sql( + s"SELECT * FROM delta.`$path` WHERE id >= (SELECT lo FROM delta.`$thresholds`) " + + s"AND id < (SELECT hi FROM delta.`$thresholds`)") + checkDeltaNativeScanAnswer(query) + + val df = query + assert(df.collect().length == 100) + // The thresholds table inside the subquery is also claimed natively; pick the + // main data-table scan by its output. + assertSubqueryFilterPushed(df, dataColumn = "v") + val scans = deltaNativeScans(df).filter(_.output.exists(_.name == "v")) + assert(scans.size == 1) + val metrics = scans.head.metrics + val rowGroupsPruned = metrics.get("row_groups_pruned_statistics").map(_.value).getOrElse(0L) + val pagesPruned = metrics.get("page_index_rows_pruned").map(_.value).getOrElse(0L) + assert( + rowGroupsPruned > 0, + s"expected row-group pruning from the resolved subquery bounds; metrics: ${metrics.map { + case (k, v) => s"$k=${v.value}" + }}") + assert( + pagesPruned > 0, + s"expected page-index pruning from the resolved subquery bounds; metrics: ${metrics.map { + case (k, v) => s"$k=${v.value}" + }}") + } + } + + test("deletion vectors: scalar subquery filter composes with DV application") { + withTempPath { dir => + val path = s"${dir.getAbsolutePath}/data" + val thresholds = s"${dir.getAbsolutePath}/thresholds" + createDvTable(path, rows = 10000) + spark.sql(s"DELETE FROM delta.`$path` WHERE id % 2 = 0") + spark + .sql("SELECT CAST(5000 AS BIGINT) AS lo") + .write + .format("delta") + .save(thresholds) + + def query = + spark.sql(s"SELECT * FROM delta.`$path` WHERE id >= (SELECT lo FROM delta.`$thresholds`)") + checkDeltaNativeScanAnswer(query) + // Deleted rows must stay deleted with the pushed bound applied in-scan. + val df = query + val rows = df.collect() + assert(rows.length == 2500) + assert(rows.forall(r => r.getLong(0) % 2 == 1 && r.getLong(0) >= 5000)) + assertSubqueryFilterPushed(df, dataColumn = "v") + } + } + + test("column mapping: scalar subquery filter on a renamed column") { + withTempPath { dir => + val path = s"${dir.getAbsolutePath}/data" + val thresholds = s"${dir.getAbsolutePath}/thresholds" + spark.range(0, 1000).selectExpr("id", "id * 2 as v").write.format("delta").save(path) + enableColumnMapping(path) + spark.sql(s"ALTER TABLE delta.`$path` RENAME COLUMN v TO w") + spark + .sql("SELECT CAST(900 AS BIGINT) AS lo") + .write + .format("delta") + .save(thresholds) + + // The pushed filter references the renamed column: it must bind against the + // physical read schema, not the logical name. + def query = + spark.sql(s"SELECT * FROM delta.`$path` WHERE w >= (SELECT lo FROM delta.`$thresholds`)") + checkDeltaNativeScanAnswer(query) + val df = query + assert(df.collect().length == 550) + assertSubqueryFilterPushed(df, dataColumn = "w") + } + } + + /** + * Assert the resolved scalar-subquery bound was actually appended to the native scan's + * execution-time common data (answers alone cannot show this: Spark's covering FilterExec would + * mask a silently-skipped pushdown). `df` must already have been executed. + */ + private def assertSubqueryFilterPushed(df: DataFrame, dataColumn: String): Unit = { + val scans = deltaNativeScans(df).collect { + case s: CometDeltaNativeScanExec if s.output.exists(_.name == dataColumn) => s + } + assert(scans.size == 1) + val scan = scans.head + val planTimeFilters = + DeltaSparkScanEnvelope.unpack(scan.nativeOp).getCommon.getDataFiltersCount + val executedFilters = OperatorOuterClass.DeltaSparkScan + .parseFrom(scan.commonData) + .getCommon + .getDataFiltersCount + assert( + executedFilters > planTimeFilters, + "expected resolved subquery filters appended at execution: " + + s"plan-time=$planTimeFilters executed=$executedFilters " + + s"dataFilters=${scan.dataFilters.mkString("; ")}") + } + + test("scalar subquery filter is NOT pushed below a limit") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 3).selectExpr("id").write.format("delta").save(path) + spark.read.format("delta").load(path).createOrReplaceTempView("t_limit_pushdown") + + val df = spark.sql( + "SELECT id FROM (SELECT id FROM t_limit_pushdown ORDER BY id LIMIT 1) q " + + "WHERE id > (SELECT max(id) FROM range(1))") + checkSparkAnswer(df) + assert(df.collect().isEmpty) + assertNoSubqueryFilterPushed(df) + } + } + + test("scalar subquery filter is NOT pushed across a nondeterministic projection") { + withSQLConf(CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.key -> "true") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 5).coalesce(1).write.format("delta").save(path) + spark.read.format("delta").load(path).createOrReplaceTempView("t_monotonic_id") + + // A deterministic conjunct does not commute with a nondeterministic projection: the + // subquery bound must not be pushed into the scan below `seq`, or the surviving rows' + // monotonically_increasing_id() values change and the answer is wrong. + val df = spark.sql( + "SELECT id FROM (SELECT id, monotonically_increasing_id() AS seq " + + "FROM t_monotonic_id) q WHERE id > (SELECT max(id) FROM range(1)) AND seq = 1") + checkSparkAnswer(df) + assert(df.collect().toSeq == Seq(Row(1))) + assertNoSubqueryFilterPushed(df) + } + } + } + + /** + * Assert no scalar-subquery filter was harvested and pushed into the native scan's + * execution-time common data: the scan must sit below a non-commuting operator (e.g. LIMIT / + * TopN), so the covering FilterExec's predicate must stay above it rather than move into the + * scan. Also confirms the query still engaged the native Delta scan, i.e. this exercises the + * commutativity guard rather than a plan that fell back to Spark entirely. `df` must already + * have been executed. + */ + private def assertNoSubqueryFilterPushed(df: DataFrame): Unit = { + val scans = deltaNativeScans(df).collect { case s: CometDeltaNativeScanExec => s } + assert(scans.size == 1, s"expected exactly one native Delta scan; found ${scans.size}") + val scan = scans.head + val planTimeFilters = + DeltaSparkScanEnvelope.unpack(scan.nativeOp).getCommon.getDataFiltersCount + val executedFilters = OperatorOuterClass.DeltaSparkScan + .parseFrom(scan.commonData) + .getCommon + .getDataFiltersCount + assert( + executedFilters == planTimeFilters, + "expected no subquery filter pushed across the non-commuting operator between the " + + s"covering filter and the scan: plan-time=$planTimeFilters executed=$executedFilters " + + s"dataFilters=${scan.dataFilters.mkString("; ")}") + } + + test("scalar subquery filter rejected by serde still marks the scan as filtered") { + withTempPath { dir => + val path = s"${dir.getAbsolutePath}/data" + val bounds = s"${dir.getAbsolutePath}/bounds" + spark.range(0, 100).selectExpr("id", "id * 2 as v").write.format("delta").save(path) + spark.sql("SELECT CAST(42 AS BIGINT) AS lo").write.format("delta").save(bounds) + + // With EqualNullSafe disabled the resolved bound cannot serialize, yet the scan must still + // carry has_data_filters so native treats it as a filtered read, exactly like core does. + withSQLConf("spark.comet.expression.EqualNullSafe.enabled" -> "false") { + def query = + spark.sql( + s"SELECT * FROM delta.`$path` WHERE id <=> (SELECT max(lo) FROM delta.`$bounds`)") + checkDeltaNativeScanAnswer(query) + val df = query + assert(df.collect().toSeq == Seq(Row(42L, 84L))) + assertUnserializedSubqueryFilterMarksScanFiltered(df, dataColumn = "v") + } + } + } + + test("unserializable scalar subquery filter keeps the safe TIMESTAMP_MILLIS conversion") { + // Same fixture as core's "filtered TIMESTAMP_MILLIS scans do not convert values Spark can + // skip": a raw file whose only overflowing millisecond value Spark prunes from the footer + // statistics once the resolved bound is pushed, so native must not convert it either. + withTempPath { dir => + val path = s"${dir.getAbsolutePath}/data" + val bounds = s"${dir.getAbsolutePath}/bounds" + writeRawParquetFile( + path, + """message root { + | optional int32 id; + | optional int64 ts(TIMESTAMP_MILLIS); + |}""".stripMargin) { factory => + (1 to 16).map(id => factory.newGroup().append("id", id).append("ts", 1717243200000L)) :+ + factory.newGroup().append("id", 17).append("ts", 9223372036854776L) + } + spark.sql(s"CONVERT TO DELTA parquet.`$path` NO STATISTICS") + spark.sql("SELECT timestamp_seconds(0) AS bound").write.format("delta").save(bounds) + + withSQLConf( + "spark.comet.expression.EqualNullSafe.enabled" -> "false", + "spark.sql.parquet.datetimeRebaseModeInRead" -> "CORRECTED", + "spark.sql.parquet.int96RebaseModeInRead" -> "CORRECTED") { + def query = spark.sql( + s"SELECT id, ts FROM delta.`$path` " + + s"WHERE ts <=> (SELECT max(bound) FROM delta.`$bounds`)") + // Spark 3.x never pushes subquery filters into its parquet reader and converts the + // overflowing value itself, so the answer comparison is meaningful on Spark 4.0+ only. + if (isSpark40Plus) { + checkDeltaNativeScanAnswer(query) + } + val df = query + assert(df.collect().isEmpty) + assert( + deltaNativeScans(df).nonEmpty, + s"expected a native Delta scan:\n${df.queryExecution}") + assertUnserializedSubqueryFilterMarksScanFiltered(df, dataColumn = "ts") + } + } + } + + /** + * Assert the execution-time common data of the scan producing `dataColumn` reports + * `has_data_filters` with no serialized data filter: the plan-time proto carries neither, and + * the resolved subquery filter is the only data filter, so only the execution-time path can set + * the bit. `df` must already have been executed. + */ + private def assertUnserializedSubqueryFilterMarksScanFiltered( + df: DataFrame, + dataColumn: String): Unit = { + val scans = deltaNativeScans(df).collect { + case s: CometDeltaNativeScanExec if s.output.exists(_.name == dataColumn) => s + } + assert(scans.size == 1, s"expected exactly one native Delta scan; found ${scans.size}") + val scan = scans.head + assert( + scan.dataFilters.exists(_.exists(_.isInstanceOf[ScalarSubquery])), + s"expected a scalar subquery data filter: ${scan.dataFilters.mkString("; ")}") + val planTime = DeltaSparkScanEnvelope.unpack(scan.nativeOp).getCommon + assert(!planTime.getHasDataFilters && planTime.getDataFiltersCount == 0) + val executed = OperatorOuterClass.DeltaSparkScan.parseFrom(scan.commonData).getCommon + assert( + executed.getHasDataFilters, + "expected has_data_filters at execution even though the resolved subquery filter did " + + s"not serialize: dataFilters=${scan.dataFilters.mkString("; ")}") + assert(executed.getDataFiltersCount == 0) + } + + test("aggregation over delta table") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark + .range(0, 10000) + .selectExpr("id", "id % 13 as g", "id * 2 as v") + .write + .format("delta") + .save(path) + + val df = spark.read + .format("delta") + .load(path) + .groupBy("g") + .sum("v") + checkDeltaNativeScanAnswer(df) + } + } + + test("conf disables the native delta scan") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 100).write.format("delta").save(path) + + withSQLConf(DeltaScanConf.COMET_DELTA_NATIVE_ENABLED.key -> "false") { + val df = spark.read.format("delta").load(path) + checkSparkAnswer(df) + assert(deltaNativeScans(df).isEmpty) + } + } + } + + test("native delta scan is opt-in: disabled when the conf is not set") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 100).write.format("delta").save(path) + + // The suite base enables the scan globally; drop the key entirely to + // observe the out-of-the-box default. + spark.conf.unset(DeltaScanConf.COMET_DELTA_NATIVE_ENABLED.key) + try { + assert(!DeltaScanConf.scanEnabled) + val df = spark.read.format("delta").load(path) + checkSparkAnswer(df) + assert(deltaNativeScans(df).isEmpty) + } finally { + spark.conf.set(DeltaScanConf.COMET_DELTA_NATIVE_ENABLED.key, "true") + } + } + } + + test("table root under a directory whose name contains a newline falls back to Spark") { + // object_store recognizes the `file` scheme but rejects the control character in the + // directory name (`%0A` in the URI), so native execution could not open the table where + // Spark's Hadoop-backed reader can. The claim gate must decline before native planning. + withTempPath { dir => + val path = new File(new File(dir, "dir\n"), "data").getAbsolutePath + spark.range(0, 100).write.format("delta").save(path) + + val df = spark.read.format("delta").load(path) + assert( + deltaNativeScans(df).isEmpty, + "Expected no native Delta scan under a newline directory:\n" + + s"${df.queryExecution.executedPlan}") + checkSparkAnswerAndFallbackReason( + df, + "Native Delta scan cannot open path 'file:" + dir.getAbsolutePath + + "/dir%0A/data': object_store rejects it") + } + } + + test("shallow clone whose source data files sit under a newline directory falls back") { + // The clone's own root is an ordinary path, so only the selected data files (resolved to + // the source table's directory) carry the rejected segment: this exercises the + // selected-paths probe, not the root gate. The reason names the first such complete path, + // a data file under the source directory. + withTempPath { dir => + val sourcePath = new File(new File(dir, "dir\n"), "source").getAbsolutePath + val clonePath = new File(dir, "clone").getAbsolutePath + spark.range(0, 100).write.format("delta").save(sourcePath) + spark.sql(s"CREATE TABLE delta.`$clonePath` SHALLOW CLONE delta.`$sourcePath`") + + val df = spark.read.format("delta").load(clonePath) + assert( + deltaNativeScans(df).isEmpty, + "Expected no native Delta scan for a clone of a newline-directory source:\n" + + s"${df.queryExecution.executedPlan}") + checkSparkAnswerAndFallbackReason( + df, + "Native Delta scan cannot open path 'file:" + dir.getAbsolutePath + + "/dir%0A/source/") + } + } + + test("converted Parquet table with a newline in a data file basename falls back to Spark") { + // CONVERT TO DELTA keeps the existing Parquet file names, so the rejected character sits in + // the basename rather than a directory segment: the table root and every parent directory + // pass the path probe, and only a check of the complete selected path can decline. + withTempPath { dir => + val path = new File(dir, "data").getAbsolutePath + spark.range(0, 100).repartition(2).write.parquet(path) + val original = new File(path).listFiles().filter(_.getName.endsWith(".parquet")).head + val renamed = new File(path, "part-00000\n.snappy.parquet") + java.nio.file.Files.move(original.toPath, renamed.toPath) + spark.sql(s"CONVERT TO DELTA parquet.`$path`") + + val df = spark.read.format("delta").load(path) + assert( + deltaNativeScans(df).isEmpty, + "Expected no native Delta scan for a converted table with a newline basename:\n" + + s"${df.queryExecution.executedPlan}") + checkSparkAnswerAndFallbackReason( + df, + s"Native Delta scan cannot open path 'file:$path/part-00000%0A.snappy.parquet': " + + "object_store rejects it") + } + } + + private def createDvTable(path: String, rows: Long = 1000): Unit = { + spark.range(0, rows).selectExpr("id", "id * 2 as v").write.format("delta").save(path) + spark.sql( + s"ALTER TABLE delta.`$path` SET TBLPROPERTIES ('delta.enableDeletionVectors' = 'true')") + } + + /** + * Same shape as `createDvTable`, plus one extra TINYINT column (value 7) under `columnName`. + */ + private def createDvTableWithExtraColumn( + path: String, + columnName: String, + rows: Long = 1000): Unit = { + spark + .range(0, rows) + .selectExpr("id", s"cast(7 as tinyint) as `$columnName`") + .write + .format("delta") + .save(path) + spark.sql( + s"ALTER TABLE delta.`$path` SET TBLPROPERTIES ('delta.enableDeletionVectors' = 'true')") + } + + test("deletion vectors: DELETE-produced DVs read natively with correct results") { + withTempPath { dir => + val path = dir.getAbsolutePath + createDvTable(path) + spark.sql(s"DELETE FROM delta.`$path` WHERE id % 2 = 0") + + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + assert(df.count() == 500) + } + } + + test( + "deletion vectors: user column named like the synthetic internal-column slot keeps its " + + "own values") { + withTempPath { dir => + val path = dir.getAbsolutePath + val collidingName = "_comet_delta___delta_internal_is_row_deleted" + createDvTableWithExtraColumn(path, collidingName) + spark.sql(s"DELETE FROM delta.`$path` WHERE id = 0") + + val df = spark.read.format("delta").load(path).select("id", collidingName) + checkDeltaNativeScanAnswer(df) + val survivingValues = df.collect().map(_.getAs[Byte](collidingName)).distinct + assert( + survivingValues.sameElements(Array(7.toByte)), + "expected the user column's own value (7) to survive DV filtering, " + + s"got ${survivingValues.toSeq}") + } + } + + test("deletion vectors: normally named extra column alongside DVs reads natively") { + withTempPath { dir => + val path = dir.getAbsolutePath + createDvTableWithExtraColumn(path, "tag") + spark.sql(s"DELETE FROM delta.`$path` WHERE id = 0") + + val df = spark.read.format("delta").load(path).select("id", "tag") + checkDeltaNativeScanAnswer(df) + val survivingValues = df.collect().map(_.getAs[Byte]("tag")).distinct + assert( + survivingValues.sameElements(Array(7.toByte)), + "expected the extra column's value (7) to survive DV filtering, got " + + survivingValues.toSeq) + } + } + + test("deletion vectors: UPDATE-produced DVs read natively with correct results") { + withTempPath { dir => + val path = dir.getAbsolutePath + createDvTable(path) + spark.sql(s"UPDATE delta.`$path` SET v = -1 WHERE id < 100") + + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + assert(df.filter(col("v") === -1).count() == 100) + assert(df.count() == 1000) + } + } + + test("deletion vectors: multiple DELETEs accumulate correctly") { + withTempPath { dir => + val path = dir.getAbsolutePath + createDvTable(path) + spark.sql(s"DELETE FROM delta.`$path` WHERE id % 2 = 0") + spark.sql(s"DELETE FROM delta.`$path` WHERE id % 3 = 0") + + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + // odd ids not divisible by 3 + assert(df.count() == (0L until 1000L).count(i => i % 2 != 0 && i % 3 != 0)) + } + } + + test( + "deletion vectors: maxDeletedRowsPerFile budget declines an oversized DV and " + + "claims once raised") { + withTempPath { dir => + val path = dir.getAbsolutePath + // repartition(4) guarantees >= 2 physical files so the per-file cardinality gate has + // more than one file to inspect. + spark + .range(0, 1000) + .selectExpr("id", "id * 2 as v") + .repartition(4) + .write + .format("delta") + .save(path) + spark.sql( + s"ALTER TABLE delta.`$path` SET TBLPROPERTIES ('delta.enableDeletionVectors' = 'true')") + spark.sql(s"DELETE FROM delta.`$path` WHERE id % 2 = 0") + + withSQLConf(DeltaScanConf.COMET_DELTA_MAX_DELETED_ROWS_PER_FILE.key -> "1") { + val df = spark.read.format("delta").load(path) + checkSparkAnswer(df) + assert( + deltaNativeScans(df).isEmpty, + "a budget of 1 deleted row per file must decline every DV-bearing file") + } + + withSQLConf(DeltaScanConf.COMET_DELTA_MAX_DELETED_ROWS_PER_FILE.key -> "1000000") { + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + } + } + } + + test("deletion vectors: maxDeletedRowsPerFile decline reason names the conf key") { + withTempPath { dir => + val path = dir.getAbsolutePath + createDvTable(path) + spark.sql(s"DELETE FROM delta.`$path` WHERE id % 2 = 0") + + withSQLConf(DeltaScanConf.COMET_DELTA_MAX_DELETED_ROWS_PER_FILE.key -> "1") { + checkSparkAnswerAndFallbackReason( + spark.read.format("delta").load(path), + DeltaScanConf.COMET_DELTA_MAX_DELETED_ROWS_PER_FILE.key) + } + } + } + + test("deletion vectors: fully-deleted region and selective predicate still prune pages") { + withTempPath { dir => + val path = dir.getAbsolutePath + val hadoopConf = spark.sparkContext.hadoopConfiguration + val oldBlockSize = hadoopConf.get("parquet.block.size") + val oldPageSize = hadoopConf.get("parquet.page.size") + hadoopConf.setInt("parquet.block.size", 256 * 1024) + hadoopConf.setInt("parquet.page.size", 16 * 1024) + try { + spark + .range(0, 500000) + .selectExpr("id", "id * 2 as v") + .sort("id") + .coalesce(1) + .write + .format("delta") + .save(path) + } finally { + if (oldBlockSize == null) hadoopConf.unset("parquet.block.size") + else hadoopConf.set("parquet.block.size", oldBlockSize) + if (oldPageSize == null) hadoopConf.unset("parquet.page.size") + else hadoopConf.set("parquet.page.size", oldPageSize) + } + spark.sql( + s"ALTER TABLE delta.`$path` SET TBLPROPERTIES ('delta.enableDeletionVectors' = 'true')") + // Delete a slice inside the predicate range and a large slice outside it. + spark.sql(s"DELETE FROM delta.`$path` WHERE id >= 150 AND id < 160") + spark.sql(s"DELETE FROM delta.`$path` WHERE id >= 300000") + + def query = spark.read + .format("delta") + .load(path) + .filter(col("id") >= 100 && col("id") < 200) + checkDeltaNativeScanAnswer(query) + + val df = query + assert(df.collect().length == 90) + val scans = deltaNativeScans(df) + assert(scans.size == 1) + val metrics = scans.head.metrics + val pagesPruned = metrics.get("page_index_rows_pruned").map(_.value).getOrElse(0L) + assert( + pagesPruned > 0, + s"expected page-index pruning to compose with DVs; metrics: ${metrics.map { case (k, v) => + s"$k=${v.value}" + }}") + } + } + + test("deletion vectors: one file split into many byte ranges reads natively") { + withTempPath { dir => + val path = dir.getAbsolutePath + // One file made of many small row groups, so a small maxPartitionBytes below turns it + // into many splits that each cover a few row groups. + val hadoopConf = spark.sparkContext.hadoopConfiguration + val oldBlockSize = hadoopConf.get("parquet.block.size") + val oldPageSize = hadoopConf.get("parquet.page.size") + hadoopConf.setInt("parquet.block.size", 16 * 1024) + hadoopConf.setInt("parquet.page.size", 4 * 1024) + try { + spark + .range(0, 20000) + .selectExpr("id", "id * 2 as v") + .coalesce(1) + .write + .format("delta") + .save(path) + } finally { + if (oldBlockSize == null) hadoopConf.unset("parquet.block.size") + else hadoopConf.set("parquet.block.size", oldBlockSize) + if (oldPageSize == null) hadoopConf.unset("parquet.page.size") + else hadoopConf.set("parquet.page.size", oldPageSize) + } + spark.sql( + s"ALTER TABLE delta.`$path` SET TBLPROPERTIES ('delta.enableDeletionVectors' = 'true')") + // Scattered deletes across every row group plus a deleted tail, so every split sees the + // deletion vector and the last splits are mostly deleted. + spark.sql(s"DELETE FROM delta.`$path` WHERE id % 5 = 0") + spark.sql(s"DELETE FROM delta.`$path` WHERE id >= 19000") + + // A claimed scan is split like any other file, and the deletion vector is applied in + // file coordinates, so each split must skip exactly its own deleted rows. + withSQLConf(SQLConf.FILES_MAX_PARTITION_BYTES.key -> "4096") { + def query = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(query) + + val df = query + assert(df.collect().length == 15200) + val scans = deltaNativeScans(df) + assert(scans.size == 1) + val numFiles = scans.head.metrics.get("numFiles").map(_.value).getOrElse(0L) + assert(numFiles == 1, s"expected a single data file, got $numFiles") + val numPartitions = scans.head.outputPartitioning.numPartitions + assert( + numPartitions > 1, + "expected the single file to be split into more than one native partition, " + + s"got $numPartitions") + } + } + } + + test("deletion vectors: aggregation over DV table") { + withTempPath { dir => + val path = dir.getAbsolutePath + createDvTable(path, rows = 10000) + spark.sql(s"DELETE FROM delta.`$path` WHERE id % 7 = 0") + + val df = spark.read.format("delta").load(path).groupBy(col("id") % 13).count() + checkDeltaNativeScanAnswer(df) + } + } + + test("deletion vectors: partitioned table reads natively with correct results") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark + .range(0, 1000) + .selectExpr("id", "id % 5 as p", "id * 2 as v") + .write + .format("delta") + .partitionBy("p") + .save(path) + spark.sql( + s"ALTER TABLE delta.`$path` SET TBLPROPERTIES ('delta.enableDeletionVectors' = 'true')") + spark.sql(s"DELETE FROM delta.`$path` WHERE id % 3 = 0") + + val df = spark.read.format("delta").load(path).filter(col("p") === 2) + checkDeltaNativeScanAnswer(df) + assert(df.count() == (0L until 1000L).count(i => i % 5 == 2 && i % 3 != 0)) + + val all = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(all) + assert(all.count() == (0L until 1000L).count(_ % 3 != 0)) + } + } + + test("deletion vectors: combined with constant metadata columns") { + withTempPath { dir => + val path = dir.getAbsolutePath + createDvTable(path) + spark.sql(s"DELETE FROM delta.`$path` WHERE id < 250") + + val df = spark.read + .format("delta") + .load(path) + .selectExpr("id", "v", "_metadata.file_name as fn") + checkSparkAnswer(df.selectExpr("id", "v", "length(fn) > 0")) + // Whether this claims or declines, results must match; if it claimed, verify the + // native node is present so the combination is actually exercised when supported. + val rows = df.collect() + assert(rows.length == 750) + assert(rows.forall(_.getString(2).nonEmpty)) + } + } + + test( + "deletion vectors: constant-metadata field names are deduplicated against the physical " + + "data and partition schemas") { + // End-to-end coverage is not possible here: selecting any `_metadata.*` field in the DV + // shape always declines today for an unrelated, pre-existing reason -- Spark reuses the + // scan's own row-index bookkeeping attribute as `_metadata.row_index`'s source, and + // `DeltaScanSupport.rowIndexUnusedAbove` conservatively treats extracting ANY `_metadata` + // field as making that attribute live (see "combined with constant metadata columns" + // above, which hedges its assertions for the same reason). That decline fires before + // `buildDvScanCommon` ever runs, regardless of collision, so it cannot exercise the fix. + // Test the builder's dedup logic directly instead, the same way `storeUris` and + // `mergedObjectStoreOptions` are unit-tested without a live scan. + val physicalDataSchema = + StructType(Seq(StructField("_comet_metadata_file_path", ByteType))) + val physicalPartitionSchema = + StructType(Seq(StructField("_comet_metadata_file_size", LongType))) + val fileConstantMetadataColumns = Seq( + AttributeReference("file_path", StringType, nullable = false)(), + AttributeReference("file_size", LongType, nullable = false)()) + + val constantMetadataFields = CometNativeScan.uniqueConstantMetadataFields( + fileConstantMetadataColumns, + physicalDataSchema.fields.map(_.name).toSet ++ physicalPartitionSchema.fields + .map(_.name) + .toSet) + assert( + constantMetadataFields.map(_.name) == Seq( + "_comet_metadata_file_path_", + "_comet_metadata_file_size_"), + "expected both constant-metadata names to be uniquified on collision, got " + + s"${constantMetadataFields.map(_.name)}") + + // The DV builder must feed these already-unique names into allocateUniqueInternalFields's + // reserved set so the internal-column suffix chain stays consistent with them. + val requiredSchema = StructType( + Seq( + StructField("id", LongType), + StructField(CometDeltaNativeScan.IsRowDeletedColumn, ByteType), + StructField(CometDeltaNativeScan.RowIndexColumn, LongType))) + val internalFields = CometDeltaNativeScan.allocateUniqueInternalFields( + requiredSchema, + physicalDataSchema, + physicalPartitionSchema, + constantMetadataFields) + + val allNames = physicalDataSchema.fields.map(_.name) ++ + physicalPartitionSchema.fields.map(_.name) ++ + constantMetadataFields.map(_.name) ++ + internalFields.map(_.name) + assert(allNames.distinct.length == allNames.length, s"expected all names distinct: $allNames") + } + + test( + "non-DV shape: user column named like the synthetic constant-metadata slot keeps its " + + "own values") { + withTempPath { dir => + val path = dir.getAbsolutePath + val collidingName = "_comet_metadata_file_path" + spark + .range(0, 100) + .selectExpr("id", s"cast(7 as tinyint) as `$collidingName`") + .write + .format("delta") + .save(path) + + val df = spark.read + .format("delta") + .load(path) + .selectExpr("id", s"`$collidingName`", "_metadata.file_path as fp") + checkDeltaNativeScanAnswer(df) + val rows = df.collect() + val survivingValues = rows.map(_.getAs[Byte](collidingName)).distinct + assert( + survivingValues.sameElements(Array(7.toByte)), + "expected the user column's own value (7) to survive the constant-metadata " + + s"collision, got ${survivingValues.toSeq}") + assert( + rows.forall(_.getString(2).nonEmpty), + "expected _metadata.file_path to still report a real path") + } + } + + test("deletion vectors: special characters in table path") { + withTempDir { base => + val dir = new java.io.File(base, "s p a r k %dv% test") + val path = dir.getAbsolutePath + createDvTable(path) + spark.sql(s"DELETE FROM delta.`$path` WHERE id % 2 = 0") + + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + assert(df.count() == 500) + } + } + + test("deletion vectors: decline when row_index is consumed via multi-hop aliases") { + withTempPath { dir => + val path = dir.getAbsolutePath + createDvTable(path) + spark.sql(s"DELETE FROM delta.`$path` WHERE id % 2 = 0") + + val df = spark.read + .format("delta") + .load(path) + .selectExpr("id", "_metadata.row_index as ri") + .selectExpr("id", "ri + 1 as ri2") + .filter(col("ri2") > 10) + checkSparkAnswer(df) + assert(deltaNativeScans(df).isEmpty, "derived row_index consumption must decline") + } + } + + test("deletion vectors: decline when row_index feeds a non-Project operator") { + withTempPath { dir => + val path = dir.getAbsolutePath + createDvTable(path) + spark.sql(s"DELETE FROM delta.`$path` WHERE id % 2 = 0") + + val df = spark.read + .format("delta") + .load(path) + .groupBy(col("_metadata.row_index") % 7) + .count() + checkSparkAnswer(df) + assert(deltaNativeScans(df).isEmpty, "aggregate over row_index must decline") + } + } + + test("deletion vectors: decline when _metadata.row_index is referenced above the scan") { + withTempPath { dir => + val path = dir.getAbsolutePath + createDvTable(path) + spark.sql(s"DELETE FROM delta.`$path` WHERE id % 2 = 0") + + val df = spark.read + .format("delta") + .load(path) + .selectExpr("id", "_metadata.row_index as ri") + checkSparkAnswer(df) + assert( + deltaNativeScans(df).isEmpty, + "plans consuming a real row_index must fall back to Spark") + } + } + + /** + * Every [[SparkPlan]] executed during `body`, captured via a [[QueryExecutionListener]] rather + * than a returned `DataFrame`'s own plan: a `DataFrameWriter` action such as `.write.parquet` + * has no result `Dataset` to call `.queryExecution` on, so the write's physical plan -- the one + * `DeltaScanSupport.declineReason` actually saw -- is only observable this way. + */ + private def capturePlansDuring(body: => Unit): Seq[SparkPlan] = { + val plans = ListBuffer.empty[SparkPlan] + val listener = new QueryExecutionListener { + override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = { + plans += qe.executedPlan + } + override def onFailure( + funcName: String, + qe: QueryExecution, + exception: Exception): Unit = {} + } + spark.listenerManager.register(listener) + try { + body + // The listener bus delivers asynchronously, so the plans are not all in hand until it has + // drained. + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + } finally { + spark.listenerManager.unregister(listener) + } + plans.toSeq + } + + test( + "deletion vectors: a write sink persisting _metadata.row_index declines the native scan " + + "and saves the real row indexes") { + withTempPath { srcDir => + withTempPath { dstDir => + val src = srcDir.getAbsolutePath + val dst = dstDir.getAbsolutePath + spark + .range(32) + .coalesce(1) + .write + .format("delta") + .option("delta.enableDeletionVectors", "true") + .save(src) + spark.sql(s"DELETE FROM delta.`$src` WHERE id IN (1, 7, 13)").collect() + + val capturedPlans = capturePlansDuring { + spark.read + .format("delta") + .load(src) + .selectExpr("id", "_metadata.row_index AS ri") + .write + .parquet(dst) + } + + // The write persists whatever the reader returns for `ri`, so the native DV scan must + // not be claimed here: claiming it would let the reader's dead synthetic row-index + // constant (correct only because the value is normally proven unused) get persisted as + // if it were the real row index. + val nativeScans = capturedPlans.flatMap(p => collectByName(p, "CometDeltaNativeScanExec")) + assert( + nativeScans.isEmpty, + "expected the write to decline the native Delta scan for a persisted row_index") + + // Documents which write-sink shape this test actually covers: the liveness gate in + // `DeltaScanSupport.rowIndexUnusedAbove` (the `childOutputLeak` check) declines a DV + // scan under ANY one-child, empty-output write sink structurally, including the DSv2 + // `V2TableWriteExec` family -- but a DV-enabled Delta read cannot be composed with a + // genuine DSv2 `AppendData` write in this delta-spark/Spark combination (see the + // dsv2-infeasibility test below), so `.write.parquet` here is the only write-sink shape + // this liveness gate is exercised against end-to-end. + assert( + capturedPlans.map(stripAQEPlan).forall(!_.isInstanceOf[V2TableWriteExec]), + s"expected a V1 write command, not a DSv2 write, got: $capturedPlans") + + val declinedScans = capturedPlans.flatMap { plan => + collectWithSubqueries(stripAQEPlan(plan)) { + case f: FileSourceScanExec if DeltaScanSupport.isDeltaScan(f) => f + } + } + val reasons = declinedScans.flatMap(f => new ExtendedExplainInfo().getFallbackReasons(f)) + assert( + reasons.exists(_.contains("row_index values consumed by the query")), + "expected the row-index-consumed-by-the-query decline reason, got: " + + reasons.mkString(", ")) + + val readBack = spark.read.parquet(dst) + checkSparkAnswer(readBack) + val rows = readBack.collect() + assert(rows.length == 29, s"expected 29 surviving rows, got ${rows.length}") + val id31 = rows.find(_.getLong(0) == 31) + assert(id31.isDefined, "expected id=31 to survive the DELETE") + assert( + id31.get.getLong(1) == 31, + "expected the persisted row_index for id=31 to be 31, got " + + s"${id31.get.getLong(1)} -- a wrongly-claimed native scan would have written a " + + "synthetic zero instead") + val sumRi = rows.map(_.getLong(1)).sum + assert( + sumRi == 475, + "expected sum(row_index) == 475 (sum(0..31) - (1 + 7 + 13) = 496 - 21), got " + + s"$sumRi -- a wrongly-claimed native scan would have summed to 0") + } + } + } + + /** + * The write-sink liveness gate above (`rowIndexUnusedAbove`'s `childOutputLeak` check in + * `DeltaScanSupport`) covers a DSv2 write sink STRUCTURALLY -- any one-child node with an empty + * output that doesn't re-expose a tainted attribute, which is exactly the shape + * `AppendDataExec`/`OverwriteByExpressionExec`/the rest of the `V2TableWriteExec` family take + * -- but the test above only ever exercises the V1 `.write.parquet` command path. + * + * Reaching a genuine DSv2 `AppendDataExec` in this Spark 3.5 setup is itself achievable: a + * table created via the session catalog with `USING parquet` still plans as a V1 + * `InsertIntoHadoopFsRelationCommand` (built-in file-based sources stay on + * `spark.sql.sources.useV1SourceList` by default), but `InMemoryTableCatalog` (from + * `spark-catalyst`'s test-jar, already a test dependency of this module, registered ad hoc + * under a throwaway name exactly as Spark's own DataSourceV2 test suites do) forces a genuine + * V2 write. + * + * What is NOT achievable in this delta-spark 3.3.2 / Spark 3.5.9 combination: composing that + * DSv2 `AppendData` write with a deletion-vector-enabled Delta table as its SOURCE. Both + * `df.writeTo(target).append()` (gluing an already-analyzed `DataFrame` into a fresh V2 + * command) AND a single `INSERT INTO target SELECT ... FROM delta.\`path\`` statement + * (resolving the read and the V2 write in one analysis pass) hit the identical failure: + * delta-spark's own `PreprocessTableWithDVs` rule requires the source relation's + * `TahoeFileIndex` to be a "pinned" `TahoeLogFileIndex` + * (`ScanWithDeletionVectors$.dvEnabledScanFor`, `PreprocessTableWithDVs.scala:78`), which does + * not hold when that relation sits under a DSv2 `AppendData` command's analysis -- confirmed + * unrelated to catalog choice or DataFrame-vs-SQL construction. This is a delta-spark + * limitation on how a DV read may be composed, not a Comet regression, so this test pins it + * down as an expected, named failure rather than silently having no DSv2 coverage at all: the + * write-sink liveness gate's DSv2 coverage for a DV row-index source remains V1-only (see the + * test above), which this test documents by construction. + */ + test( + "deletion vectors: a genuine DSv2 AppendData write cannot compose with a DV-enabled Delta " + + "source in this Spark/Delta combination (delta-spark's own pinned-snapshot requirement, " + + "not a Comet regression) -- documents why DSv2 write-sink coverage stays V1-only above") { + val catalogName = "cometDeltaRowIndexV2Cat" + withSQLConf( + s"spark.sql.catalog.$catalogName" -> + "org.apache.spark.sql.connector.catalog.InMemoryTableCatalog") { + withTempPath { srcDir => + val src = srcDir.getAbsolutePath + spark + .range(32) + .coalesce(1) + .write + .format("delta") + .option("delta.enableDeletionVectors", "true") + .save(src) + spark.sql(s"DELETE FROM delta.`$src` WHERE id IN (1, 7, 13)").collect() + + val targetTable = s"$catalogName.ns.row_index_sink" + spark.sql(s"CREATE TABLE $targetTable (id BIGINT, ri BIGINT) USING foo") + + val ex = intercept[IllegalArgumentException] { + spark.sql( + s"INSERT INTO $targetTable SELECT id, _metadata.row_index AS ri FROM delta.`$src`") + } + assert( + ex.getMessage.contains("non-pinned"), + "expected delta-spark's pinned-TahoeLogFileIndex requirement to be the failure " + + "(if this now succeeds, DSv2 coverage for the DV row-index write-sink scenario " + + "may finally be achievable and this test should be replaced with a real one): " + + ex.getMessage) + } + } + } + + /** + * Single-file (ids 0-4) deletion-vector table with one id deleted, used by the UnionExec + * row-index liveness tests below: UnionExec's output takes its expression IDs positionally from + * its FIRST child, so a live `_metadata.row_index` alias in a later branch is invisible to a + * taint analysis that only follows `ProjectExec` aliases. A small fixed fixture keeps the + * expected surviving row_index values easy to hand-verify. + */ + private def createSmallDvTable(path: String, deleteId: Long): Unit = { + spark.range(0, 5).selectExpr("id").coalesce(1).write.format("delta").save(path) + spark.sql( + s"ALTER TABLE delta.`$path` SET TBLPROPERTIES ('delta.enableDeletionVectors' = 'true')") + spark.sql(s"DELETE FROM delta.`$path` WHERE id = $deleteId") + } + + test( + "deletion vectors: row_index live through UNION ALL declines both branches " + + "with correct SUM") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withTempPath { dir1 => + withTempPath { dir2 => + val t1 = dir1.getAbsolutePath + val t2 = dir2.getAbsolutePath + // t1: ids 0,1,3,4 survive (id 2 deleted); t2: ids 0,1,2,4 survive (id 3 deleted). + createSmallDvTable(t1, deleteId = 2) + createSmallDvTable(t2, deleteId = 3) + + def query: DataFrame = { + val left = + spark.read.format("delta").load(t1).selectExpr("id", "_metadata.row_index as ri") + val right = + spark.read.format("delta").load(t2).selectExpr("id", "_metadata.row_index as ri") + left.union(right) + } + + checkSparkAnswer(query) + val df = query + val rows = df.collect() + assert(rows.length == 8, s"expected 8 surviving rows, got ${rows.length}") + assert( + deltaNativeScans(df).isEmpty, + "row_index live via a union's positional output remap must decline both branches") + // row_index equals id for every surviving row in this single-file, insertion-ordered + // fixture, so summing the real (uncorrupted) row indexes is equivalent to summing ids: + // t1 (0+1+3+4=8) + t2 (0+1+2+4=7) = 15. A wrongly-claimed branch would instead + // contribute a constant 0 per row, which this exact total rules out. + val sum = rows.map(_.getLong(1)).sum + assert(sum == 15L, s"expected SUM(ri) == 15, got $sum") + } + } + } + } + + test( + "deletion vectors: row_index live only in the second UNION ALL branch declines " + + "only that branch") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withTempPath { dir1 => + withTempPath { dir2 => + val t1 = dir1.getAbsolutePath + val t2 = dir2.getAbsolutePath + createSmallDvTable(t1, deleteId = 2) + createSmallDvTable(t2, deleteId = 3) + + def query: DataFrame = { + // Branch 1's "ri" is a constant, never derived from its own row_index; branch 2's + // "ri" is the real _metadata.row_index. UnionExec's output reuses branch 1's + // expression ID for the "ri" column, so only branch 2's scan should decline. + val left = + spark.read.format("delta").load(t1).selectExpr("id", "CAST(-1 AS BIGINT) as ri") + val right = + spark.read.format("delta").load(t2).selectExpr("id", "_metadata.row_index as ri") + left.union(right) + } + + checkSparkAnswer(query) + val df = query + val rows = df.collect() + assert(rows.length == 8, s"expected 8 surviving rows, got ${rows.length}") + val fromT1 = rows.filter(_.getLong(1) == -1L) + val fromT2 = rows.filter(_.getLong(1) != -1L) + assert(fromT1.length == 4, s"expected 4 rows from t1, got ${fromT1.length}") + assert(fromT2.length == 4, s"expected 4 rows from t2, got ${fromT2.length}") + // Real row_index equals id in this fixture; a wrongly-claimed branch 2 would instead + // report a constant 0 for every row, which this per-row check rules out. + assert( + fromT2.forall(r => r.getLong(0) == r.getLong(1)), + s"expected t2's ri to equal id, got: ${fromT2.mkString(", ")}") + val scans = deltaNativeScans(df) + assert( + scans.size == 1, + s"expected exactly branch 1 (t1) to claim natively, got ${scans.size} native scans") + } + } + } + } + + test( + "deletion vectors: SUM(row_index) over UNION ALL declines both branches with correct total") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withTempPath { dir1 => + withTempPath { dir2 => + val t1 = dir1.getAbsolutePath + val t2 = dir2.getAbsolutePath + createSmallDvTable(t1, deleteId = 2) + createSmallDvTable(t2, deleteId = 3) + + def query: DataFrame = { + val left = + spark.read.format("delta").load(t1).selectExpr("id", "_metadata.row_index as ri") + val right = + spark.read.format("delta").load(t2).selectExpr("id", "_metadata.row_index as ri") + left.union(right).selectExpr("sum(ri) as total") + } + + checkSparkAnswer(query) + val total = query.collect()(0).getLong(0) + assert(total == 15L, s"expected SUM(ri) == 15, got $total") + assert( + deltaNativeScans(query).isEmpty, + "row_index live via an aggregate over a union must decline both branches") + } + } + } + } + + test( + "deletion vectors: UNION ALL without _metadata still claims both branches natively " + + "(anti-regression)") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withTempPath { dir1 => + withTempPath { dir2 => + val t1 = dir1.getAbsolutePath + val t2 = dir2.getAbsolutePath + createSmallDvTable(t1, deleteId = 2) + createSmallDvTable(t2, deleteId = 3) + + def query: DataFrame = { + val left = spark.read.format("delta").load(t1).selectExpr("id") + val right = spark.read.format("delta").load(t2).selectExpr("id") + left.union(right) + } + + checkSparkAnswer(query) + val df = query + val ids = df.collect().map(_.getLong(0)).sorted + assert( + ids.sameElements(Array(0L, 0L, 1L, 1L, 2L, 3L, 4L, 4L)), + s"unexpected surviving ids: ${ids.mkString(", ")}") + val scans = deltaNativeScans(df) + assert( + scans.size == 2, + "a DV union without _metadata must still claim both branches natively " + + s"(the row-index column is dead in both), got ${scans.size} native scans") + } + } + } + } + + test("deletion vectors: inner join between two DV tables claims both scans natively") { + // Positive coverage for the generic multi-child safety net (DeltaScanSupport.scala's + // multiChildLeak check): a plain join carries no row-index taint at all, so the safety net + // must not mistake a join's normal attribute passthrough for a leak and fall both sides back + // to Spark. + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withTempPath { dir1 => + withTempPath { dir2 => + val t1 = dir1.getAbsolutePath + val t2 = dir2.getAbsolutePath + // t1 survives ids {0,1,3,4} (id 2 deleted); t2 survives ids {0,1,2,4} (id 3 deleted). + createSmallDvTable(t1, deleteId = 2) + createSmallDvTable(t2, deleteId = 3) + + def query: DataFrame = { + val left = spark.read.format("delta").load(t1).withColumnRenamed("id", "lid") + val right = spark.read.format("delta").load(t2).withColumnRenamed("id", "rid") + left.join(right, col("lid") === col("rid")) + } + + checkSparkAnswer(query) + val df = query + val rows = df.collect() + val ids = rows.map(_.getLong(0)).sorted + // Only ids surviving in BOTH tables' deletion vectors should match. + assert( + ids.sameElements(Array(0L, 1L, 4L)), + s"expected join to match surviving ids {0,1,4}, got: ${ids.mkString(", ")}") + assert( + rows.forall(r => r.getLong(0) == r.getLong(1)), + "join key mismatch in result rows") + val scans = deltaNativeScans(df) + assert( + scans.size == 2, + "a DV-backed join with no row-index consumption must claim both sides natively, " + + s"got ${scans.size} native scans") + } + } + } + } + + private def enableColumnMapping(path: String): Unit = + spark.sql(s"""ALTER TABLE delta.`$path` SET TBLPROPERTIES ( + | 'delta.minReaderVersion' = '2', + | 'delta.minWriterVersion' = '5', + | 'delta.columnMapping.mode' = 'name')""".stripMargin) + + test("column mapping: renamed column reads natively across old and new files") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 100).selectExpr("id", "id * 2 as v").write.format("delta").save(path) + enableColumnMapping(path) + spark.sql(s"ALTER TABLE delta.`$path` RENAME COLUMN v TO w") + // Files written after the rename carry the same physical name. + spark + .range(100, 200) + .selectExpr("id", "id * 2 as w") + .write + .format("delta") + .mode("append") + .save(path) + + val df = spark.read.format("delta").load(path).filter(col("w") > 100) + checkDeltaNativeScanAnswer(df) + assert(spark.read.format("delta").load(path).count() == 200) + } + } + + test("column mapping: dropped and re-added column name reads natively") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 100).selectExpr("id", "id * 2 as v").write.format("delta").save(path) + enableColumnMapping(path) + spark.sql(s"ALTER TABLE delta.`$path` DROP COLUMN v") + spark.sql(s"ALTER TABLE delta.`$path` ADD COLUMN v LONG") + spark + .range(100, 200) + .selectExpr("id", "id * 3 as v") + .write + .format("delta") + .mode("append") + .save(path) + + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + // Old files must yield NULL for the re-added v (different physical column). + assert(df.filter(col("id") < 100).filter(col("v").isNotNull).count() == 0) + assert(df.filter(col("id") >= 100).filter(col("v").isNull).count() == 0) + } + } + + test("column mapping: partitioned table reads natively") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark + .range(0, 500) + .selectExpr("id", "id % 5 as p") + .write + .format("delta") + .partitionBy("p") + .save(path) + enableColumnMapping(path) + spark.sql(s"ALTER TABLE delta.`$path` RENAME COLUMN p TO part") + + val df = spark.read.format("delta").load(path).filter(col("part") === 3) + checkDeltaNativeScanAnswer(df) + assert(df.count() == 100) + } + } + + test( + "column mapping: rename history colliding logical partition name with physical data " + + "name reads correctly") { + withTempPath { dir => + val path = dir.getAbsolutePath + // a->b then p->a leaves the LOGICAL name "a" bound to the partition column while a + // DIFFERENT physical data column (originally "a", now logically "b") retains physical + // name "a". Passing the partition schema's logical names to the native side collides + // with that retained physical data name and lets DataFusion's name-based partition + // rewrite replace the data projection with the partition constant. + spark.sql(s"CREATE TABLE delta.`$path` (a BIGINT, p BIGINT) USING delta PARTITIONED BY (p)") + enableColumnMapping(path) + spark.sql(s"INSERT INTO delta.`$path` VALUES (1, 100), (2, 100)") + spark.sql(s"ALTER TABLE delta.`$path` RENAME COLUMN a TO b") + spark.sql(s"ALTER TABLE delta.`$path` RENAME COLUMN p TO a") + + val df = spark.sql(s"SELECT b, a FROM delta.`$path`") + checkDeltaNativeScanAnswer(df) + val rows = df.collect().map(r => (r.getLong(0), r.getLong(1))).sorted + assert( + rows.sameElements(Array((1L, 100L), (2L, 100L))), + s"expected (1,100),(2,100) but got ${rows.mkString(", ")}") + } + } + + test( + "column mapping: rename history colliding partition name reads correctly with " + + "deletion vectors") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.sql(s"CREATE TABLE delta.`$path` (a BIGINT, p BIGINT) USING delta PARTITIONED BY (p)") + enableColumnMapping(path) + spark.sql( + s"ALTER TABLE delta.`$path` SET TBLPROPERTIES ('delta.enableDeletionVectors' = 'true')") + spark.sql(s"INSERT INTO delta.`$path` VALUES (1, 100), (2, 100)") + spark.sql(s"ALTER TABLE delta.`$path` RENAME COLUMN a TO b") + spark.sql(s"ALTER TABLE delta.`$path` RENAME COLUMN p TO a") + spark.sql(s"DELETE FROM delta.`$path` WHERE b = 1") + + val df = spark.sql(s"SELECT b, a FROM delta.`$path`") + checkDeltaNativeScanAnswer(df) + val rows = df.collect().map(r => (r.getLong(0), r.getLong(1))).sorted + assert( + rows.sameElements(Array((2L, 100L))), + s"expected (2,100) but got ${rows.mkString(", ")}") + } + } + + test("column mapping: renamed partition column without collision reads correctly (control)") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.sql(s"CREATE TABLE delta.`$path` (a BIGINT, p BIGINT) USING delta PARTITIONED BY (p)") + enableColumnMapping(path) + spark.sql(s"INSERT INTO delta.`$path` VALUES (1, 100), (2, 100)") + // Rename ONLY the partition column, to a name that collides with nothing: no physical + // data column is named "q", so this must not be affected by the collision above. + spark.sql(s"ALTER TABLE delta.`$path` RENAME COLUMN p TO q") + + val df = spark.sql(s"SELECT a, q FROM delta.`$path`") + checkDeltaNativeScanAnswer(df) + val rows = df.collect().map(r => (r.getLong(0), r.getLong(1))).sorted + assert( + rows.sameElements(Array((1L, 100L), (2L, 100L))), + s"expected (1,100),(2,100) but got ${rows.mkString(", ")}") + } + } + + test("column mapping: combined with deletion vectors") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 1000).selectExpr("id", "id * 2 as v").write.format("delta").save(path) + enableColumnMapping(path) + spark.sql( + s"ALTER TABLE delta.`$path` SET TBLPROPERTIES ('delta.enableDeletionVectors' = 'true')") + spark.sql(s"ALTER TABLE delta.`$path` RENAME COLUMN v TO w") + spark.sql(s"DELETE FROM delta.`$path` WHERE id % 4 = 0") + + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + assert(df.count() == 750) + } + } + + test("column mapping: to_json on a nested struct matches Spark's field names") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark + .range(0, 10) + .selectExpr("id", "named_struct('a', id) as s") + .write + .format("delta") + .save(path) + enableColumnMapping(path) + // Renaming the NESTED field (not the outer column) is what diverges the physical name + // ("a", preserved on rename) from the logical name ("b") for a struct field below the + // top level -- the shape that leaks physical names into to_json's output. + spark.sql(s"ALTER TABLE delta.`$path` RENAME COLUMN s.a TO b") + + withSQLConf(CometConf.getExprAllowIncompatConfigKey(classOf[StructsToJson]) -> "true") { + val df = spark.read.format("delta").load(path).select(to_json(col("s"))) + checkSparkAnswer(df) + assert( + deltaNativeScans(df).isEmpty, + "column mapping with nested struct fields must fall back to Spark") + } + } + } + + test("decline: column mapping with nested struct columns falls back to Spark") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark + .range(0, 200) + .selectExpr( + "id", + "named_struct('a', id, 'b', cast(id as string)) as st", + "array(id, id * 2) as arr", + "map(cast(id as string), id) as mp") + .write + .format("delta") + .save(path) + enableColumnMapping(path) + spark.sql(s"ALTER TABLE delta.`$path` RENAME COLUMN st TO st2") + spark + .range(200, 300) + .selectExpr( + "id", + "named_struct('a', id, 'b', cast(id as string)) as st2", + "array(id, id * 2) as arr", + "map(cast(id as string), id) as mp") + .write + .format("delta") + .mode("append") + .save(path) + + val df = spark.read.format("delta").load(path).selectExpr("id", "st2.a", "arr", "mp") + checkSparkAnswer(df) + assert(df.count() == 300) + assert( + deltaNativeScans(df).isEmpty, + "column mapping with nested struct fields must fall back to Spark") + } + } + + test("decline: column mapping with structs nested in arrays and maps falls back to Spark") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark + .range(0, 50) + .selectExpr( + "id", + "array(named_struct('a', id, 'b', cast(id as string))) as arrOfStruct", + "map(cast(id as string), named_struct('a', id)) as mapOfStruct") + .write + .format("delta") + .save(path) + enableColumnMapping(path) + + val df = spark.read.format("delta").load(path) + checkSparkAnswer(df) + assert( + deltaNativeScans(df).isEmpty, + "column mapping with structs nested in arrays/maps must fall back to Spark") + } + } + + test("column mapping: top-level scalars and array-of-primitives still claim natively") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark + .range(0, 200) + .selectExpr("id", "cast(id as string) as v", "array(id, id * 2) as arr") + .write + .format("delta") + .save(path) + enableColumnMapping(path) + spark.sql(s"ALTER TABLE delta.`$path` RENAME COLUMN v TO w") + + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + assert(df.count() == 200) + } + } + + test("decline: column mapping id mode falls back to Spark with correct results") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.sql(s"""CREATE TABLE delta.`$path` (id LONG, v LONG) USING delta + |TBLPROPERTIES ('delta.columnMapping.mode' = 'id')""".stripMargin) + spark + .range(0, 100) + .selectExpr("id", "id * 2 as v") + .write + .format("delta") + .mode("append") + .save(path) + + val df = spark.read.format("delta").load(path) + checkSparkAnswer(df) + assert(deltaNativeScans(df).isEmpty, "id-mode column mapping must decline") + } + } + + test("delete without deletion vectors rewrites files and still reads natively") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 1000).selectExpr("id", "id * 2 as v").write.format("delta").save(path) + // DVs are off by default, so DELETE rewrites files; result is still a plain table. + spark.sql(s"DELETE FROM delta.`$path` WHERE id < 100") + + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + assert(df.count() == 900) + } + } + + test("dynamic partition pruning via broadcast join prunes delta partitions") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "20", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "100m") { + withTempPath { factDir => + withTempPath { dimDir => + val factPath = factDir.getAbsolutePath + val dimPath = dimDir.getAbsolutePath + spark + .range(0, 2000) + .selectExpr("id", "id % 10 as p") + .write + .format("delta") + .partitionBy("p") + .save(factPath) + // Unfiltered on disk: the selective predicate below is a query-time filter, which + // is what gives Spark's DynamicPartitionPruning rule a subquery to inject in the + // first place. Filtering before the write (the old shape of this test) leaves no + // predicate in the query for DPP to see, so the assertions below never fired. + spark + .range(0, 10) + .selectExpr("id as key", "id % 10 as dp") + .write + .format("delta") + .save(dimPath) + + def query = { + val fact = spark.read.format("delta").load(factPath) + val dim = spark.read.format("delta").load(dimPath) + // only partitions 0 and 1 survive the join + fact.join(dim, fact("p") === dim("dp")).filter(dim("key") < 2) + } + + checkSparkAnswer(query) + + val df = query + val rows = df.collect() + assert(rows.length == 400) // 2 partitions x 200 rows + val scans = deltaNativeScans(df) + assert( + scans.nonEmpty, + s"expected native delta scans:\n${df.queryExecution.executedPlan}") + + val deltaScans = scans.collect { case s: CometDeltaNativeScanExec => s } + assert( + deltaScans.exists(_.runtimeFilters.exists(_.isInstanceOf[DynamicPruningExpression])), + "expected a DynamicPruningExpression in a CometDeltaNativeScanExec's " + + s"runtimeFilters:\n${df.queryExecution.executedPlan}") + + // The fact-side scan must have read fewer files than the table holds (DPP pruning). + val factScan = scans.maxBy(_.metrics.get("staticFilesNum").map(_.value).getOrElse(0L)) + val staticFiles = factScan.metrics.get("staticFilesNum").map(_.value).getOrElse(0L) + val readFiles = factScan.metrics.get("numFiles").map(_.value).getOrElse(0L) + assert(staticFiles > 0, "expected the staticFilesNum metric to be populated") + assert( + readFiles < staticFiles, + s"expected DPP pruning: read $readFiles of $staticFiles files") + } + } + } + } + + for (adaptive <- Seq(false, true)) { + test( + "DPP stays executable when transition reversion restores the delta scan: " + + s"AQE=$adaptive") { + assume(isSpark35Plus, "Comet AQE DPP query-stage optimizer rules require Spark 3.5+") + withTempPath { factDir => + withTempPath { dimDir => + val factPath = factDir.getAbsolutePath + val dimPath = dimDir.getAbsolutePath + spark + .range(0, 400) + .selectExpr( + "CAST(id AS INT) AS fact_id", + "CAST(id % 10 AS INT) AS fact_key", + "CONCAT('f', id) AS fact_str") + .write + .format("delta") + .partitionBy("fact_key") + .save(factPath) + spark + .range(0, 10) + .selectExpr( + "CAST(id AS INT) AS dim_id", + "CAST(id AS INT) AS dim_key", + "CONCAT('d', id) AS dim_str") + .write + .parquet(dimPath) + + withTempView("revert_dpp_fact", "revert_dpp_dim") { + spark.read.format("delta").load(factPath).createOrReplaceTempView("revert_dpp_fact") + spark.read.parquet(dimPath).createOrReplaceTempView("revert_dpp_dim") + val query = + """SELECT f.fact_id, f.fact_str, d.dim_str + |FROM revert_dpp_fact f JOIN revert_dpp_dim d ON f.fact_key = d.dim_key + |WHERE d.dim_id < 5""".stripMargin + + def hasDpp(filters: Seq[Expression]): Boolean = + filters.exists(_.isInstanceOf[DynamicPruningExpression]) + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", + "spark.comet.exec.project.enabled" -> "false", + CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") { + // Without reversion the fact side is a native Delta scan carrying the DPP filter. + withSQLConf(CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "false") { + val df = sql(query) + assert(df.collect().length == 200) + val nativeScans = deltaNativeScans(df).collect { + case s: CometDeltaNativeScanExec => s + } + assert( + nativeScans.exists(s => hasDpp(s.runtimeFilters)), + "expected a native Delta scan with a DPP filter:\n" + + df.queryExecution.executedPlan) + } + + withSQLConf(CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true") { + val (_, cometPlan) = checkSparkAnswer(query) + assert( + collectByName(cometPlan, "CometDeltaNativeScanExec").isEmpty, + s"transition reversion should restore the Spark Delta scan:\n$cometPlan") + val restored = collect(cometPlan) { + case s: FileSourceScanExec if hasDpp(s.partitionFilters) => s + } + assert( + restored.nonEmpty, + s"expected a restored scan with a DPP filter:\n$cometPlan") + restored.foreach { scan => + assert( + scan.relation.fileFormat.getClass.getSimpleName.contains("Delta"), + s"expected the restored scan to read Delta:\n$cometPlan") + assert( + scan.logicalLink.isDefined, + s"restored scan lost its logical link:\n$scan") + val subqueryPlans = + scan.partitionFilters.flatMap(_.collect { case e: InSubqueryExec => + e.plan match { + case ReusedSubqueryExec(p) => p + case p => p + } + }) + assert( + subqueryPlans.nonEmpty && subqueryPlans.forall { + case _: SubqueryBroadcastExec | _: CometSubqueryBroadcastExec => true + case _ => false + }, + s"restored scan must carry the executable DPP subquery:\n$cometPlan") + } + } + } + } + } + } + } + } + + test("union all with DPP join and coalescible shuffle survives AQE partitioning checks") { + // The crash shape: a DPP join in one UNION ALL branch and a coalescible shuffle (the + // GROUP BY) in the other. Spark's AQE plan validation walks every operator's + // outputPartitioning, including the DPP branch's scan, before + // CometPlanAdaptiveDynamicPruningFilters has rewritten the placeholder subquery -- this + // is the ordering that reproduced the crash. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "20", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "100m") { + withTempPath { factDir => + withTempPath { dimDir => + withTempPath { otherDir => + val factPath = factDir.getAbsolutePath + val dimPath = dimDir.getAbsolutePath + val otherPath = otherDir.getAbsolutePath + + spark + .range(0, 2000) + .selectExpr("id", "id % 10 as p") + .write + .format("delta") + .partitionBy("p") + .save(factPath) + spark + .range(0, 10) + .selectExpr("id as key", "id % 10 as dp", "id as sel") + .write + .format("delta") + .save(dimPath) + spark + .range(0, 500) + .selectExpr("id", "id % 10 as p") + .write + .format("delta") + .partitionBy("p") + .save(otherPath) + + spark.read.format("delta").load(factPath).createOrReplaceTempView("r43Fact") + spark.read.format("delta").load(dimPath).createOrReplaceTempView("r43Dim") + spark.read.format("delta").load(otherPath).createOrReplaceTempView("r43Other") + + def query = + spark.sql(""" + |SELECT f.p, f.id FROM r43Fact f JOIN r43Dim d ON f.p = d.dp WHERE d.sel < 2 + |UNION ALL + |SELECT p, CAST(count(*) AS LONG) AS id FROM r43Other GROUP BY p + |""".stripMargin) + + try { + checkSparkAnswer(query) + } catch { + case e: Throwable => + if (e.getMessage != null && + e.getMessage.contains("does not support the execute() code path")) { + throw new AssertionError( + "AQE inspected outputPartitioning on an unresolved adaptive DPP " + + "placeholder -- this is the crash this test guards against", + e) + } + throw e + } + + val df = query + df.collect() + // Best-effort: this UNION ALL shape need not always route through the native + // Delta scan, but if it does, it must have survived AQE's partitioning checks + // above without throwing. Observed to vary run-to-run on this build (Spark 3.5.9 + // / Delta 3.3.2), so this is logged rather than asserted -- answer correctness is + // already verified by checkSparkAnswer above. + val scans = deltaNativeScans(df) + if (scans.isEmpty) { + logInfo( + "union all with DPP join and coalescible shuffle: no CometDeltaNativeScanExec " + + "claimed this query on this build; answer correctness already verified above") + } else { + logInfo( + s"union all with DPP join and coalescible shuffle: ${scans.length} " + + "CometDeltaNativeScanExec node(s) claimed this query; answer correctness " + + "already verified above") + } + } + } + } + } + } + + test("scalar subquery in a partition filter does not force partitioning during AQE checks") { + // Crash shape: a scalar subquery used directly as + // a partition filter, e.g. `p = (SELECT max(p) FROM dim ...)`, references only the + // partition column, so it lands in runtimeFilters rather than dataFilters. + // ValidateRequirements walks outputPartitioning for every operator, including this scan, + // before the subquery has executed -- forcing perPartitionData at that point evaluates the + // still-unresolved ScalarSubquery and throws "has not finished". + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "20", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "100m") { + withTempPath { factDir => + withTempPath { dimDir => + withTempPath { otherDir => + val factPath = factDir.getAbsolutePath + val dimPath = dimDir.getAbsolutePath + val otherPath = otherDir.getAbsolutePath + + spark + .range(0, 2000) + .selectExpr("id", "id % 10 as p") + .write + .format("delta") + .partitionBy("p") + .save(factPath) + spark + .range(0, 10) + .selectExpr("id as p", "case when id in (0, 3) then 'yes' else 'no' end as country") + .write + .format("parquet") + .save(dimPath) + spark + .range(0, 500) + .selectExpr("id", "id % 10 as p") + .write + .format("parquet") + .save(otherPath) + + spark.read.format("delta").load(factPath).createOrReplaceTempView("r45Fact") + spark.read.format("parquet").load(dimPath).createOrReplaceTempView("r45Dim") + spark.read.format("parquet").load(otherPath).createOrReplaceTempView("r45Other") + + def query = + spark.sql(""" + |SELECT id, p FROM r45Fact + |WHERE p = (SELECT max(p) FROM r45Dim WHERE country = 'yes') + |UNION ALL + |SELECT cast(count(*) AS int) AS id, p FROM r45Other GROUP BY p + |""".stripMargin) + + try { + checkSparkAnswer(query) + } catch { + case e: Throwable => + if (e.getMessage != null && e.getMessage.contains("has not finished")) { + throw new AssertionError( + "AQE ValidateRequirements forced outputPartitioning to evaluate an " + + "unresolved scalar partition-filter subquery -- this is the crash this " + + "test guards against", + e) + } + throw e + } + + val df = query + df.collect() + // Best-effort, mirroring the DPP union-all test above: this shape need not always + // route through the native Delta scan, but if it does, it must have survived AQE's + // partitioning checks above without throwing. Answer correctness is already + // verified by checkSparkAnswer above. + val scans = deltaNativeScans(df) + if (scans.isEmpty) { + logInfo( + "scalar subquery partition filter: no CometDeltaNativeScanExec claimed this " + + "query on this build; answer correctness already verified above") + } else { + logInfo( + s"scalar subquery partition filter: ${scans.length} CometDeltaNativeScanExec " + + "node(s) claimed this query; answer correctness already verified above") + } + } + } + } + } + } + + test( + "aggregate over a scalar-subquery partition filter executes under a fused native " + + "parent") { + // Crash shape: a scalar subquery used as a partition filter (`p = (SELECT max(p) ...)`) + // lands in runtimeFilters. Once execution resolves it, a native aggregate sitting + // directly on top of the scan (no intervening exchange) reads the scan's + // outputPartitioning to size its own execution context; that getter must report the + // real post-pruning partition count, not a value stuck from before resolution. + withTempPath { factDir => + withTempPath { thresholdsDir => + val factPath = factDir.getAbsolutePath + val thresholdsPath = thresholdsDir.getAbsolutePath + + spark + .range(0, 2000) + .selectExpr("id", "id % 10 as p") + .write + .format("delta") + .partitionBy("p") + .save(factPath) + + spark + .sql("SELECT CAST(7 AS BIGINT) AS p") + .write + .format("delta") + .save(thresholdsPath) + + def query = + spark.sql( + s"SELECT sum(id) AS total FROM delta.`$factPath` " + + s"WHERE p = (SELECT max(p) FROM delta.`$thresholdsPath`)") + + checkSparkAnswer(query) + + val df = query + try { + df.collect() + } catch { + case e: Throwable => + if (e.getMessage != null && e.getMessage.contains("All per-partition arrays")) { + throw new AssertionError( + "a fused native aggregate above the scan read a stale zero " + + "outputPartitioning after the scalar-subquery partition filter had " + + "already resolved", + e) + } + throw e + } + + val scans = deltaNativeScans(df) + assert( + scans.nonEmpty, + s"expected CometDeltaNativeScanExec in plan:\n${df.queryExecution.executedPlan}") + assert( + collectByName(df.queryExecution.executedPlan, "CometHashAggregateExec").nonEmpty, + "expected a fused native aggregate parent above the scan in plan:\n" + + s"${df.queryExecution.executedPlan}") + } + } + } + + test( + "metrics evaluates without throwing when runtimeFilters holds a ScalarSubquery " + + "placeholder (pins the invariant documented on CometDeltaNativeScanExec.scanHelper: " + + "AQE's UI plan-walk calls .metrics on every node mid-planning, sometimes before a " + + "DPP/scalar-subquery filter has resolved, and this must never throw)") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 50).write.format("delta").save(path) + + val scan = deltaNativeScans(spark.read.format("delta").load(path)).collect { + case s: CometDeltaNativeScanExec => s + }.head + + // A real execution.ScalarSubquery instance (the exec-time class CometDeltaNativeScanExec + // itself matches against in hasUnevaluableSubqueryFilter), wrapping a never-executed + // SubqueryExec -- deliberately never run, so this is unresolved exactly as it would be + // when AQE's mid-planning walk reaches this node ahead of subquery execution. + val innerPlan = spark.range(1).selectExpr("id AS c").queryExecution.executedPlan + val unresolvedScalarSubquery = + ScalarSubquery( + SubqueryExec("metrics-guard-subquery", innerPlan), + NamedExpression.newExprId) + + val scanWithSubquery = scan.copy(runtimeFilters = Seq(unresolvedScalarSubquery)) + val metrics = scanWithSubquery.metrics + assert( + metrics.nonEmpty, + "expected CometDeltaNativeScanExec.metrics to populate the native scan metric node " + + "even with an unresolved ScalarSubquery in runtimeFilters") + } + } + + test("input_file_name falls back to Spark with correct results") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 100).selectExpr("id").write.format("delta").save(path) + + val df = spark.read + .format("delta") + .load(path) + .selectExpr("id", "input_file_name() as f") + checkSparkAnswer(df.selectExpr("id", "length(f) > 0")) + assert(deltaNativeScans(df).isEmpty, "input_file_name must decline") + } + } + + test("self-join of the same delta table keeps scans distinct") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 100).selectExpr("id", "id % 5 as k").write.format("delta").save(path) + + def query = { + val left = spark.read.format("delta").load(path).filter(col("id") < 50) + val right = spark.read.format("delta").load(path).filter(col("id") >= 50) + left.as("l").join(right.as("r"), col("l.k") === col("r.k")) + } + checkSparkAnswer(query) + + val df = query + df.collect() + assert(deltaNativeScans(df).size == 2) + } + } + + test("schema evolution: added column yields nulls for old files, natively") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 100).selectExpr("id").write.format("delta").save(path) + spark.sql(s"ALTER TABLE delta.`$path` ADD COLUMN v LONG") + spark + .range(100, 200) + .selectExpr("id", "id * 2 as v") + .write + .format("delta") + .mode("append") + .save(path) + + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + assert(df.filter(col("id") < 100).filter(col("v").isNotNull).count() == 0) + assert(df.filter(col("id") >= 100).filter(col("v").isNull).count() == 0) + } + } + + test("schema evolution: column default (Delta two-step) reads correctly") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 100).selectExpr("id").write.format("delta").save(path) + spark.sql( + s"ALTER TABLE delta.`$path` SET TBLPROPERTIES " + + "('delta.feature.allowColumnDefaults' = 'supported')") + // Delta only allows defaults via add-then-set (applies to FUTURE inserts; old files + // read as NULL -- unlike Spark's existence defaults). + spark.sql(s"ALTER TABLE delta.`$path` ADD COLUMN v LONG") + spark.sql(s"ALTER TABLE delta.`$path` ALTER COLUMN v SET DEFAULT 42") + spark.sql(s"INSERT INTO delta.`$path` (id) VALUES (100), (101)") + + val df = spark.read.format("delta").load(path) + // Whether claimed or declined, results must match Spark exactly. + checkSparkAnswer(df) + assert(df.count() == 102) + assert(df.filter(col("v") === 42).count() == 2) + assert(df.filter(col("id") < 100).filter(col("v").isNotNull).count() == 0) + } + } + + test("legacy INT96 timestamps read natively with correct values") { + withTempPath { dir => + val path = dir.getAbsolutePath + withSQLConf("spark.sql.parquet.outputTimestampType" -> "INT96") { + spark + .range(0, 100) + .selectExpr("id", "timestamp_seconds(1600000000 + id * 3600) as ts") + .write + .format("delta") + .save(path) + } + val df = spark.read.format("delta").load(path).filter(col("id") < 50) + checkDeltaNativeScanAnswer(df) + assert(df.count() == 50) + } + } + + test("decline: type widening feature falls back with correct results") { + withTempPath { dir => + val path = dir.getAbsolutePath + // Delta 3.3's widening preview supports byte/short -> int. + spark.sql(s"""CREATE TABLE delta.`$path` (id SMALLINT) USING delta + |TBLPROPERTIES ('delta.enableTypeWidening' = 'true')""".stripMargin) + spark + .range(0, 100) + .selectExpr("cast(id as smallint) as id") + .write + .format("delta") + .mode("append") + .save(path) + spark.sql(s"ALTER TABLE delta.`$path` ALTER COLUMN id TYPE INT") + spark + .range(100, 200) + .selectExpr("cast(id as int) as id") + .write + .format("delta") + .mode("append") + .save(path) + + val df = spark.read.format("delta").load(path) + checkSparkAnswer(df) + assert(df.count() == 200) + } + } + + test( + "decline: SMALLINT column falls back with correct results when unsigned-small-int " + + "safety check is enabled") { + // Regression: the Delta claim path must + // run the same CometScanTypeChecker core's own scan does, so the default-on + // COMET_PARQUET_UNSIGNED_SMALL_INT_CHECK safety fallback still applies to a native Delta + // scan. Without it, an out-of-range/malformed UINT_8 payload stored under a ShortType + // column could be claimed and silently decoded with the wrong values. + withTempPath { dir => + val path = dir.getAbsolutePath + spark.sql(s"CREATE TABLE delta.`$path` (id INT, s SMALLINT) USING delta") + spark.sql(s"INSERT INTO delta.`$path` VALUES (1, 10), (2, 20), (3, 30)") + + // CometTestBase flips this conf off by default so the rest of the suite can exercise + // ShortType columns against Comet's native scan; put it back to its real production + // default so this gate actually declines (mirrors the same pattern in + // DeltaScanContribSuite for the vectorized-reader conf). + withSQLConf(CometConf.COMET_PARQUET_UNSIGNED_SMALL_INT_CHECK.key -> "true") { + val df = spark.read.format("delta").load(path) + checkSparkAnswerAndFallbackReason( + df, + CometConf.COMET_PARQUET_UNSIGNED_SMALL_INT_CHECK.key) + assert(deltaNativeScans(df).isEmpty) + } + } + } + + test("claims SMALLINT column natively when unsigned-small-int safety check is disabled") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.sql(s"CREATE TABLE delta.`$path` (id INT, s SMALLINT) USING delta") + spark.sql(s"INSERT INTO delta.`$path` VALUES (1, 10), (2, 20), (3, 30)") + + withSQLConf(CometConf.COMET_PARQUET_UNSIGNED_SMALL_INT_CHECK.key -> "false") { + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + } + } + } + + test("checkpointed delta log reads natively") { + withTempPath { dir => + val path = dir.getAbsolutePath + // Force a checkpoint by exceeding the default interval via many commits. + spark.sql(s"""CREATE TABLE delta.`$path` (id LONG, v LONG) USING delta + |TBLPROPERTIES ('delta.checkpointInterval' = '3')""".stripMargin) + for (i <- 0 until 5) { + spark + .range(i * 10L, (i + 1) * 10L) + .selectExpr("id", "id * 2 as v") + .write + .format("delta") + .mode("append") + .save(path) + } + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + assert(df.count() == 50) + } + } + + test( + "decline: shallow clone with a supported local root but viewfs-scheme selected files " + + "falls back to Spark") { + // A shallow clone whose root is a supported scheme (local `file:`) but whose selected data + // files resolve through the source table's unsupported location (`viewfs:`, mounted onto + // the local filesystem so results are checkable). A gate on rootPaths alone misses this. + val cluster = "cometDeltaViewfsGate" + // Hadoop's mounttable is plain Configuration, not SQLConf: mutate the session's shared + // hadoopConfiguration directly (mirroring withSQLConf's set-then-restore shape) rather than + // withSQLConf, which only round-trips actual SQLConf entries. + val hadoopConf = spark.sparkContext.hadoopConfiguration + val linkFallbackKey = s"fs.viewfs.mounttable.$cluster.linkFallback" + val priorLinkFallback = Option(hadoopConf.get(linkFallbackKey)) + hadoopConf.set(linkFallbackKey, "file:///") + try { + withTempPath { sourceDir => + withTempPath { cloneDir => + val sourcePath = sourceDir.getAbsolutePath + val clonePath = cloneDir.getAbsolutePath + val sourceViewfsPath = s"viewfs://$cluster$sourcePath" + + spark + .range(0, 10) + .write + .format("delta") + .save(sourceViewfsPath) + spark.sql(s"CREATE TABLE delta.`$clonePath` SHALLOW CLONE delta.`$sourceViewfsPath`") + // Append local (file:) data on top of the clone's inherited viewfs-scheme files: the + // scan's selected data files now span both an unsupported scheme (viewfs) AND multiple + // object-store authorities (file: carries none, viewfs://cometDeltaViewfsGate carries + // one), the same shape DeltaScanContribSuite's + // "unsupportedSelectedSchemeReason declines a mixed file:+viewfs selection" unit test + // pins directly against declineReason's gate ordering (DeltaScanSupport.scala): the + // scheme gate runs before multiStoreReason, so the fallback reason below must still + // name viewfs, never "spans multiple object stores". This confirms that ordering + // end to end through declineReason, not merely at the unit level. + spark.range(10, 20).write.format("delta").mode("append").save(clonePath) + + val df = spark.read.format("delta").load(clonePath) + assert( + deltaNativeScans(df).isEmpty, + "Expected no native Delta scan for a viewfs-selected-file clone:\n" + + s"${df.queryExecution.executedPlan}") + checkSparkAnswerAndFallbackReason( + df, + "Native Delta scan does not support selected data file or deletion vector " + + "filesystem scheme(s) viewfs") + } + } + } finally { + priorLinkFallback match { + case Some(v) => hadoopConf.set(linkFallbackKey, v) + case None => hadoopConf.unset(linkFallbackKey) + } + } + } + + test("change data feed read never engages the native Delta scan, with correct results") { + // A batch readChangeFeed() query never reaches DeltaScanSupport.declineReason's own + // isCDCRead check at all: CDCReader wraps its answer in a DeltaCDFRelation whose buildScan + // executes its internal (possibly DeltaParquetFileFormat-backed) plan via queryExecution's + // RDD lineage directly, so the physical plan Spark and Comet's extensions ultimately see for + // this query is a single, opaque RowDataSourceScanExec, never a FileSourceScanExec + // DeltaScanSupport.isDeltaScan could recognize. This still pins the outcome that matters: + // Change Data Feed reads are never claimed by the native Delta scan and stay correct. + withTempPath { dir => + val path = dir.getAbsolutePath + // Change Data Feed must be enabled from the table's first version: CDC reads validate + // that change data was actually recorded for every version in the requested range. + spark.sql(s"""CREATE TABLE delta.`$path` (id LONG, v LONG) USING delta + |TBLPROPERTIES ('delta.enableChangeDataFeed' = 'true')""".stripMargin) + spark.sql(s"INSERT INTO delta.`$path` SELECT id, id * 2 FROM range(0, 100)") + spark.sql(s"UPDATE delta.`$path` SET v = -1 WHERE id < 10") + + val df = spark.read + .format("delta") + .option("readChangeFeed", "true") + .option("startingVersion", 0) + .load(path) + checkSparkAnswer(df) + assert(deltaNativeScans(df).isEmpty) + assert(df.count() > 0) + } + } + + test("reader features: TIMESTAMP_NTZ column claims natively with correct results") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.sql(s"CREATE TABLE delta.`$path` (id LONG, ts TIMESTAMP_NTZ) USING delta") + spark.sql( + s"INSERT INTO delta.`$path` VALUES " + + "(1, CAST('2021-01-01 00:00:00' AS TIMESTAMP_NTZ)), " + + "(2, CAST('2022-06-15 12:30:00' AS TIMESTAMP_NTZ))") + + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + assert(df.count() == 2) + } + } + + test("reader features: v2Checkpoint table claims natively with correct results") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.sql(s"""CREATE TABLE delta.`$path` (id LONG, v LONG) USING delta + |TBLPROPERTIES ( + | 'delta.checkpointPolicy' = 'v2', + | 'delta.checkpointInterval' = '3')""".stripMargin) + for (i <- 0 until 5) { + spark + .range(i * 10L, (i + 1) * 10L) + .selectExpr("id", "id * 2 as v") + .write + .format("delta") + .mode("append") + .save(path) + } + + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + assert(df.count() == 50) + } + } + + test( + "reader features: an unsupported reader feature (type widening) declines with the " + + "reader feature(s) reason") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.sql(s"""CREATE TABLE delta.`$path` (id SMALLINT) USING delta + |TBLPROPERTIES ('delta.enableTypeWidening' = 'true')""".stripMargin) + spark + .range(0, 100) + .selectExpr("cast(id as smallint) as id") + .write + .format("delta") + .mode("append") + .save(path) + spark.sql(s"ALTER TABLE delta.`$path` ALTER COLUMN id TYPE INT") + + val df = spark.read.format("delta").load(path) + checkSparkAnswerAndFallbackReason( + df, + "Native Delta scan does not support reader feature(s) typeWidening") + assert(deltaNativeScans(df).isEmpty) + } + } + + test("_metadata.row_index declines before any deletion vector exists on a DV-enabled table") { + // _metadata.row_index only resolves on a Delta table once deletion-vector support is on + // the protocol (it errors as an unknown field otherwise); once it resolves, Delta always + // routes the read through the DV-application shape (a row-index column with no + // is_row_deleted alongside it), even with zero deletion vectors written yet. This pins that + // the hasRowIndex-without-hasIsRowDeleted gate declines this shape regardless of whether a + // DV has ever actually been written for the file. + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 100).selectExpr("id").write.format("delta").save(path) + spark.sql( + s"ALTER TABLE delta.`$path` SET TBLPROPERTIES ('delta.enableDeletionVectors' = 'true')") + + val df = spark.read + .format("delta") + .load(path) + .selectExpr("id", "_metadata.row_index as ri") + checkSparkAnswerAndFallbackReason( + df, + "Native Delta scan does not support row-index reads outside a deletion-vector scan") + assert(deltaNativeScans(df).isEmpty) + assert(df.count() == 100) + } + } + + test( + "decline: parquet.crypto.factory.class configured declines conservatively even without " + + "actual encryption") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 100).selectExpr("id", "id * 2 as v").write.format("delta").save(path) + + val hadoopConf = spark.sparkContext.hadoopConfiguration + val key = "parquet.crypto.factory.class" + val prior = Option(hadoopConf.get(key)) + // A real, resolvable factory that explicitly allows plaintext files: the table itself is + // NOT encrypted, so this exercises Comet's stricter, conservative "decline ALL + // encrypted-parquet configurations" gate without breaking Spark's own read. + hadoopConf.set(key, "org.apache.parquet.crypto.keytools.PropertiesDrivenCryptoFactory") + try { + val df = spark.read.format("delta").load(path) + checkSparkAnswerAndFallbackReason( + df, + "Native Delta scan does not support encrypted parquet") + assert(deltaNativeScans(df).isEmpty) + assert(df.count() == 100) + } finally { + prior match { + case Some(v) => hadoopConf.set(key, v) + case None => hadoopConf.unset(key) + } + } + } + } + + test( + "deletion vectors: a data predicate deleting every row of one file still claims " + + "natively with correct results") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark + .range(0, 40) + .selectExpr("id", "id % 2 as p", "id * 2 as v") + .repartition(2, col("p")) + .write + .format("delta") + .partitionBy("p") + .save(path) + spark.sql( + s"ALTER TABLE delta.`$path` SET TBLPROPERTIES ('delta.enableDeletionVectors' = 'true')") + // A data-column predicate (not purely a partition predicate) forces Delta through the + // row-level deletion-vector path rather than a metadata-only partition drop, even though + // every row in partition 1's file happens to match. + spark.sql(s"DELETE FROM delta.`$path` WHERE p = 1 AND v >= 0") + + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + assert(df.count() == 20) + assert(df.filter(col("p") === 1).count() == 0) + } + } + + test( + "conf interactions: ANSI, case sensitivity, and disabled DPP leave claim/decline " + + "outcomes unchanged") { + withTempPath { claimDir => + withTempPath { declineDir => + val claimPath = claimDir.getAbsolutePath + val declinePath = declineDir.getAbsolutePath + spark.range(0, 200).selectExpr("id", "id * 2 as v").write.format("delta").save(claimPath) + spark.sql(s"""CREATE TABLE delta.`$declinePath` (id LONG, v LONG) USING delta + |TBLPROPERTIES ('delta.columnMapping.mode' = 'id')""".stripMargin) + spark + .range(0, 200) + .selectExpr("id", "id * 2 as v") + .write + .format("delta") + .mode("append") + .save(declinePath) + + val confVariants = Seq( + SQLConf.ANSI_ENABLED.key -> "true", + SQLConf.CASE_SENSITIVE.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "false") + + confVariants.foreach { case (key, value) => + withSQLConf(key -> value) { + val claimDf = spark.read.format("delta").load(claimPath) + checkDeltaNativeScanAnswer(claimDf) + + val declineDf = spark.read.format("delta").load(declinePath) + checkSparkAnswer(declineDf) + assert( + deltaNativeScans(declineDf).isEmpty, + s"expected id-mode column mapping to still decline under $key=$value") + } + } + } + } + } + + test( + "deletion vectors: maxDeletedRowsPerFile boundary claims when cardinality exactly " + + "equals the limit (gate declines only when the limit is exceeded)") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark + .range(0, 1000) + .selectExpr("id", "id * 2 as v") + .coalesce(1) + .write + .format("delta") + .save(path) + spark.sql( + s"ALTER TABLE delta.`$path` SET TBLPROPERTIES ('delta.enableDeletionVectors' = 'true')") + spark.sql(s"DELETE FROM delta.`$path` WHERE id % 2 = 0") + + withSQLConf(DeltaScanConf.COMET_DELTA_MAX_DELETED_ROWS_PER_FILE.key -> "500") { + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + assert(df.count() == 500) + } + } + } + + // Both tests below pin caseSensitive=true purely to exercise the exact-match (non-folding) + // path for a non-ASCII column name. Native's case-insensitive name matching reproduces the + // JVM's `toLowerCase(Locale.ROOT)` fold (see `fold_names` in + // native/core/src/parquet/name_fold.rs), so caseSensitive=false would also read these + // correctly -- there is no decline gate involved here to route around. + test("unicode column names round-trip natively with correct results") { + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "true") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.sql(s"CREATE TABLE delta.`$path` (id LONG, `名前` STRING) USING delta") + spark.sql(s"INSERT INTO delta.`$path` VALUES (1, 'たろう'), (2, 'はなこ')") + + val df = spark.sql(s"SELECT id, `名前` FROM delta.`$path` ORDER BY id") + checkDeltaNativeScanAnswer(df) + val rows = df.collect() + assert(rows.map(_.getString(1)).sameElements(Array("たろう", "はなこ"))) + } + } + } + + test("unicode and space-containing column names round-trip natively under column mapping") { + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "true") { + withTempPath { dir => + val path = dir.getAbsolutePath + // A space is one of Parquet's disallowed schema-name characters, so the space-containing + // column can only be added AFTER column mapping (physical names) is already active -- + // creating it inline at CREATE TABLE time fails before column mapping ever takes effect. + spark.sql(s"CREATE TABLE delta.`$path` (id LONG, `名前` STRING) USING delta") + enableColumnMapping(path) + spark.sql(s"ALTER TABLE delta.`$path` ADD COLUMN `a b` LONG") + spark.sql(s"INSERT INTO delta.`$path` VALUES (1, 'たろう', 10), (2, 'はなこ', 20)") + + val df = spark.sql(s"SELECT id, `名前`, `a b` FROM delta.`$path` ORDER BY id") + checkDeltaNativeScanAnswer(df) + val rows = df.collect() + assert(rows.map(_.getString(1)).sameElements(Array("たろう", "はなこ"))) + assert(rows.map(_.getLong(2)).sameElements(Array(10L, 20L))) + } + } + } + + test( + "a non-ASCII case-insensitive column name claims the native Delta scan with correct " + + "results") { + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { + withTempPath { dir => + val path = dir.getAbsolutePath + val table = "comet_unicode_" + java.util.UUID.randomUUID().toString.replace("-", "") + withTable(table) { + // Two plain parquet files whose footers differ only in the case of a non-ASCII letter + // (an ordinary CONVERT-eligible layout: no column mapping, no defaults, no DVs). + // Native's name matcher reproduces this JVM's `toLowerCase(Locale.ROOT)` from + // shipped case tables, which folds 'É'/'é' together just like Spark does. + spark.range(1, 2).select(col("id"), lit(71).as("É")).coalesce(1).write.parquet(path) + spark + .range(2, 3) + .select(col("id"), lit(72).as("é")) + .coalesce(1) + .write + .mode("append") + .parquet(path) + + spark.sql(s"CREATE TABLE $table (id BIGINT, `É` INT) USING PARQUET LOCATION '$path'") + spark.sql(s"CONVERT TO DELTA $table NO STATISTICS") + + val df = spark.read.format("delta").load(path).selectExpr("id", "`É`") + checkDeltaNativeScanAnswer(df) + val rows = df.collect().sortBy(_.getLong(0)) + assert(rows.map(_.getInt(1)).sameElements(Array(71, 72))) + } + } + } + } + + test( + "an ASCII case-insensitive column name still claims the native Delta scan with correct " + + "results") { + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { + withTempPath { dir => + val path = dir.getAbsolutePath + val table = "comet_ascii_case_" + java.util.UUID.randomUUID().toString.replace("-", "") + withTable(table) { + // Same shape as above, but the differing-case letter is plain ASCII, which native's + // name folding (`fold_names` in name_fold.rs) always matches correctly, ASCII being + // the easy case. + spark.range(1, 2).select(col("id"), lit(71).as("E")).coalesce(1).write.parquet(path) + spark + .range(2, 3) + .select(col("id"), lit(72).as("e")) + .coalesce(1) + .write + .mode("append") + .parquet(path) + + spark.sql(s"CREATE TABLE $table (id BIGINT, `E` INT) USING PARQUET LOCATION '$path'") + spark.sql(s"CONVERT TO DELTA $table NO STATISTICS") + + val df = spark.read.format("delta").load(path).selectExpr("id", "`E`") + checkDeltaNativeScanAnswer(df) + val rows = df.collect().sortBy(_.getLong(0)) + assert(rows.map(_.getInt(1)).sameElements(Array(71, 72))) + } + } + } + } + + test( + "a non-ASCII partition column name still claims the native Delta scan with correct " + + "results") { + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { + withTempPath { dir => + val path = dir.getAbsolutePath + // Partition values are injected into the output as constants by exact name match, never + // matched against a file's footer schema, so a non-ASCII partition name (data names stay + // plain ASCII here) never goes through native's case-insensitive DATA-column name + // folding (`fold_names` in name_fold.rs) at all. + spark + .range(0, 20) + .selectExpr("id", "cast(id % 4 as long) as `名前`") + .write + .format("delta") + .partitionBy("名前") + .save(path) + + val df = spark.read.format("delta").load(path).filter(col("名前") === 2) + checkDeltaNativeScanAnswer(df) + assert(df.count() > 0) + } + } + } + + test( + "a non-ASCII physical column name still claims the native Delta scan under column " + + "mapping with case-insensitive reads") { + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { + withTempPath { dir => + val path = dir.getAbsolutePath + // The column pre-exists the column-mapping upgrade, so Delta assigns its physical name + // as its current (non-ASCII) name verbatim -- exactly what a converted-then-upgraded + // table keeps. Logical and physical names are identical here, so this was always safe; + // it now also claims natively rather than being caught by a blanket non-ASCII gate. + spark.sql(s"CREATE TABLE delta.`$path` (id LONG, `É` STRING) USING delta") + enableColumnMapping(path) + spark.sql(s"INSERT INTO delta.`$path` VALUES (1, 'a'), (2, 'b')") + + val df = spark.sql(s"SELECT id, `É` FROM delta.`$path` ORDER BY id") + checkDeltaNativeScanAnswer(df) + val rows = df.collect() + assert(rows.map(_.getString(1)).sameElements(Array("a", "b"))) + } + } + } + + test( + "a Kelvin sign physical column name in one file of an otherwise-ASCII CONVERTed table " + + "claims the native Delta scan with correct results") { + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { + withTempPath { dir => + val path = dir.getAbsolutePath + val table = "comet_kelvin_" + java.util.UUID.randomUUID().toString.replace("-", "") + withTable(table) { + // An ordinary CONVERT-eligible layout (no column mapping, no defaults, no DVs) + // where the table is declared with a plain ASCII "K" column, but one of its + // underlying Parquet files happens to have been written with a physical column + // literally named U+212A (KELVIN SIGN) -- not decomposable to ASCII by naive + // folding, but a case variant of ASCII 'k'/'K' under Java's `Character` mappings + // (and thus under Spark's `caseSensitive=false` resolution). Nothing on the JVM + // side can see this: the divergent name lives only in the second file's footer. + spark.range(1, 2).select(col("id"), lit(71).as("K")).coalesce(1).write.parquet(path) + spark + .range(2, 3) + .select(col("id"), lit(72).as("K")) + .coalesce(1) + .write + .mode("append") + .parquet(path) + + spark.sql(s"CREATE TABLE $table (id BIGINT, `K` INT) USING PARQUET LOCATION '$path'") + spark.sql(s"CONVERT TO DELTA $table NO STATISTICS") + + val df = spark.read.format("delta").load(path).selectExpr("id", "`K`") + checkDeltaNativeScanAnswer(df) + val rows = df.collect().sortBy(_.getLong(0)) + assert(rows.map(_.getInt(1)).sameElements(Array(71, 72))) + assert( + df.filter(col("K").isNotNull).count() == 2, + "the Kelvin-sign-named file's row must not be nulled out by native") + } + } + } + } + + test( + "a capital-sigma physical column name matches a final-sigma table column with correct " + + "results") { + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { + withTempPath { dir => + val path = dir.getAbsolutePath + val table = "comet_sigma_" + java.util.UUID.randomUUID().toString.replace("-", "") + withTable(table) { + // Java's `String.toLowerCase(Locale.ROOT)` lowers "A1Σ" to "a1ς" (FINAL + // sigma): its Final_Cased context scan runs on word boundaries, and the digit keeps + // "A1Σ" a single word, so the trailing sigma takes the final form. Spark's + // footer matching therefore folds physical "A1Σ" onto a requested "a1ς", + // and the value in that file must be read, not nulled. Nothing on the JVM side can + // see this: the divergent name lives only in the second file's footer. + spark + .range(1, 2) + .select(col("id"), lit(71).as("a1ς")) + .coalesce(1) + .write + .parquet(path) + spark + .range(2, 3) + .select(col("id"), lit(72).as("A1Σ")) + .coalesce(1) + .write + .mode("append") + .parquet(path) + + spark.sql(s"CREATE TABLE $table (id BIGINT, `a1ς` INT) USING PARQUET LOCATION '$path'") + spark.sql(s"CONVERT TO DELTA $table NO STATISTICS") + + val df = spark.read.format("delta").load(path).selectExpr("id", "`a1ς`") + checkDeltaNativeScanAnswer(df) + val rows = df.collect().sortBy(_.getLong(0)) + assert(rows.map(_.getInt(1)).sameElements(Array(71, 72))) + assert( + df.filter(col("a1ς").isNotNull).count() == 2, + "the capital-sigma-named file's row must not be nulled out by native") + } + } + } + } + + test("a capital-sigma physical column name is missing for a non-final-sigma table column") { + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { + withTempPath { dir => + val path = dir.getAbsolutePath + val table = "comet_sigma_miss_" + java.util.UUID.randomUUID().toString.replace("-", "") + withTable(table) { + // The inverse of the test above: "A1Σ" lowers to "a1ς", NOT "a1σ" + // (non-final sigma), so Spark's footer lookup treats a requested "a1σ" as + // MISSING in the capital-sigma file and substitutes NULL. Reading a value there + // (as a naive codepoint-wise fold would) surfaces a row Spark considers absent + // and breaks IS NOT NULL filters. + spark + .range(1, 2) + .select(col("id"), lit(71).as("a1σ")) + .coalesce(1) + .write + .parquet(path) + spark + .range(2, 3) + .select(col("id"), lit(72).as("A1Σ")) + .coalesce(1) + .write + .mode("append") + .parquet(path) + + spark.sql(s"CREATE TABLE $table (id BIGINT, `a1σ` INT) USING PARQUET LOCATION '$path'") + spark.sql(s"CONVERT TO DELTA $table NO STATISTICS") + + val df = spark.read.format("delta").load(path).selectExpr("id", "`a1σ`") + checkDeltaNativeScanAnswer(df) + val rows = df.collect().sortBy(_.getLong(0)) + assert(rows.length == 2) + assert(rows(0).getInt(1) == 71) + assert( + rows(1).isNullAt(1), + "the capital-sigma file's column lowers to final sigma, so a non-final-sigma " + + "requested column must read as missing (NULL) there") + } + } + } + } + + test("a Unicode-version-drift physical column name folds per the running JDK") { + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { + withTempPath { dir => + val path = dir.getAbsolutePath + val table = "comet_drift_" + java.util.UUID.randomUUID().toString.replace("-", "") + withTable(table) { + // U+A7C0 (LATIN CAPITAL LETTER OLD POLISH O) gained its lowercase pairing U+A7C1 + // in Unicode 14, after JDK 17's Unicode snapshot: JDK 17 lowers it to itself + // (no match against a U+A7C1 column), while JDK 21+ lowers it to U+A7C1 (match). + // The expectation is derived from the RUNNING JDK's own toLowerCase, so this test + // is correct on any JDK -- exactly the property the native matcher must mirror, + // since it consumes case tables generated by this same JVM at plan time. + val physicalFolds = + "Ꟁ".toLowerCase(java.util.Locale.ROOT) == "ꟁ" + + spark + .range(1, 2) + .select(col("id"), lit(71).as("ꟁ")) + .coalesce(1) + .write + .parquet(path) + spark + .range(2, 3) + .select(col("id"), lit(72).as("Ꟁ")) + .coalesce(1) + .write + .mode("append") + .parquet(path) + + spark.sql(s"CREATE TABLE $table (id BIGINT, `ꟁ` INT) USING PARQUET LOCATION '$path'") + spark.sql(s"CONVERT TO DELTA $table NO STATISTICS") + + val df = spark.read.format("delta").load(path).selectExpr("id", "`ꟁ`") + checkDeltaNativeScanAnswer(df) + val rows = df.collect().sortBy(_.getLong(0)) + assert(rows.length == 2) + assert(rows(0).getInt(1) == 71) + if (physicalFolds) { + assert( + !rows(1).isNullAt(1) && rows(1).getInt(1) == 72, + "this JDK folds U+A7C0 onto U+A7C1, so the value must be read") + } else { + assert( + rows(1).isNullAt(1), + "this JDK does not fold U+A7C0 onto U+A7C1, so the column must be missing") + } + } + } + } + } + + /** + * Runs `action` under a [[SparkListener]] that captures every `onTaskEnd` input-metrics + * reading, then drains the listener bus before summing: the bus delivers `onTaskEnd` + * asynchronously, so `action` returning is not enough to guarantee every event has already been + * processed. Callers compare the sums against a floor rather than an exact target because + * Delta's own transaction-log state reconstruction runs a small auxiliary job reading the + * commit JSON, which legitimately contributes a few extra input records alongside the actual + * data scan. Returns the aggregated (recordsRead, bytesRead). + */ + private def collectTaskInputMetrics(action: => Unit): (Long, Long) = { + val inputRecords = mutable.ArrayBuffer.empty[Long] + val inputBytes = mutable.ArrayBuffer.empty[Long] + val listener = new SparkListener { + override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = { + val im = taskEnd.taskMetrics.inputMetrics + inputRecords.synchronized { inputRecords += im.recordsRead } + inputBytes.synchronized { inputBytes += im.bytesRead } + } + } + spark.sparkContext.addSparkListener(listener) + try { + action + CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + (inputRecords.synchronized(inputRecords.sum), inputBytes.synchronized(inputBytes.sum)) + } finally { + spark.sparkContext.removeSparkListener(listener) + } + } + + test("standalone uncached delta read reports task-level input metrics") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark + .range(0, 10000) + .selectExpr("id", "id * 2 as v") + .write + .format("delta") + .save(path) + + val df = spark.read.format("delta").load(path) + var collected = 0L + val (recordsRead, bytesRead) = collectTaskInputMetrics { + collected = df.collect().length.toLong + } + + assert(collected == 10000L) + assert( + deltaNativeScans(df).nonEmpty, + s"expected a native Delta scan:\n${df.queryExecution.executedPlan}") + assert( + recordsRead >= 10000L, + s"expected task input recordsRead to cover the row count, got $recordsRead") + assert(bytesRead > 0L, s"expected task input bytesRead > 0, got $bytesRead") + } + } + + test("fused aggregate over a delta scan reports task-level input metrics") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark + .range(0, 10000) + .selectExpr("id", "id % 13 as g", "id * 2 as v") + .write + .format("delta") + .save(path) + + val df = spark.read.format("delta").load(path).groupBy("g").sum("v") + val (recordsRead, bytesRead) = collectTaskInputMetrics { + df.collect() + } + + assert( + deltaNativeScans(df).nonEmpty, + s"expected the native Delta scan fused into the aggregate:\n${df.queryExecution.executedPlan}") + assert( + recordsRead >= 10000L, + s"expected task input recordsRead to cover the scanned row count, got $recordsRead") + assert(bytesRead > 0L, s"expected task input bytesRead > 0, got $bytesRead") + } + } + + /** Write `values` rows (id, d, ts) as a Delta table with the given write-side rebase modes. */ + private def writeRebaseTable( + path: String, + timeZone: String, + datetimeMode: String, + int96Mode: String, + values: String): Unit = { + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> timeZone, + "spark.sql.parquet.datetimeRebaseModeInWrite" -> datetimeMode, + "spark.sql.parquet.int96RebaseModeInWrite" -> int96Mode) { + spark + .sql(s"select * from values $values as t(id, d, ts)") + .write + .format("delta") + .save(path) + } + } + + test("legacy-rebase ancient dates and timestamps match Spark's own read") { + // Spark stamps org.apache.spark.legacyDateTime / legacyINT96 / timeZone into the file + // footer when writing with LEGACY rebase modes, and its own reader rebases based on that + // per-file metadata regardless of the session's read-mode conf. The native scan must + // resolve the same per-file policy: without rebasing, 1500-01-01 reads as 1500-01-10. + withTempPath { dir => + val path = dir.getAbsolutePath + writeRebaseTable( + path, + timeZone = "UTC", + datetimeMode = "LEGACY", + int96Mode = "LEGACY", + values = "(1, date'0001-01-01', timestamp'1500-01-01 00:00:00'), " + + "(2, date'1500-01-01', timestamp'1582-10-04 23:59:59'), " + + "(3, date'1582-10-04', timestamp'0001-01-01 00:00:00'), " + + "(4, date'2024-06-01', timestamp'2024-06-01 12:00:00')") + withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC") { + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + } + } + } + + test("predicate on a legacy-rebase ancient date matches Spark") { + withTempPath { dir => + val path = dir.getAbsolutePath + writeRebaseTable( + path, + timeZone = "UTC", + datetimeMode = "LEGACY", + int96Mode = "LEGACY", + values = "(1, date'1500-01-01', timestamp'1500-01-01 00:00:00'), " + + "(2, date'1500-02-11', timestamp'1500-02-11 00:00:00'), " + + "(3, date'2024-06-01', timestamp'2024-06-01 12:00:00')") + withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC") { + val df = spark.read + .format("delta") + .load(path) + .filter("d = date'1500-01-01'") + checkDeltaNativeScanAnswer(df) + } + } + } + + test("legacy-rebase file holding only modern values stays native and correct") { + // Rebasing is the identity from 1582-10-15 onward, so a LEGACY-stamped file whose values + // are all modern must keep reading natively with unchanged results. + withTempPath { dir => + val path = dir.getAbsolutePath + writeRebaseTable( + path, + timeZone = "UTC", + datetimeMode = "LEGACY", + int96Mode = "LEGACY", + values = "(1, date'1990-01-01', timestamp'1990-01-01 00:00:00'), " + + "(2, date'2024-06-01', timestamp'2024-06-01 12:00:00')") + withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC") { + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + } + } + } + + test( + "legacy-rebase ancient timestamps with a non-UTC writer zone fail loudly instead of " + + "returning shifted values") { + // Timestamp rebasing outside a fixed UTC writer zone needs the JVM's historical timezone + // tables; the native reader refuses ancient values rather than guessing. + withTempPath { dir => + val path = dir.getAbsolutePath + writeRebaseTable( + path, + timeZone = "America/Los_Angeles", + datetimeMode = "LEGACY", + int96Mode = "LEGACY", + values = "(1, date'2024-06-01', timestamp'1500-01-01 00:00:00')") + withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> "America/Los_Angeles") { + val e = intercept[Exception] { + spark.read.format("delta").load(path).collect() + } + val messages = Iterator + .iterate(e: Throwable)(_.getCause) + .takeWhile(_ != null) + .map(_.getMessage) + .mkString("\n") + assert(messages.contains("rebase"), s"expected a calendar-rebase error, got:\n$messages") + } + } + } + + test("mixed rebase flags attribute each timestamp column to its physical type's flag") { + // legacyDateTime governs INT64 timestamps while legacyINT96 governs INT96 ones. A file + // carrying exactly one of the two flags must read every timestamp column under the flag + // of its own physical type -- rebased exactly when that flag is LEGACY, verbatim when it + // is not -- matching Spark's own read, instead of refusing ancient values because the two + // flags disagree. All four (physical type, mode pair) combinations round-trip + // 1500-01-01 00:00:00. + for ((outputType, datetimeMode, int96Mode) <- Seq( + ("TIMESTAMP_MICROS", "LEGACY", "CORRECTED"), + ("TIMESTAMP_MICROS", "CORRECTED", "LEGACY"), + ("INT96", "LEGACY", "CORRECTED"), + ("INT96", "CORRECTED", "LEGACY"))) { + withTempPath { dir => + val path = dir.getAbsolutePath + withSQLConf("spark.sql.parquet.outputTimestampType" -> outputType) { + writeRebaseTable( + path, + timeZone = "UTC", + datetimeMode = datetimeMode, + int96Mode = int96Mode, + values = "(1, date'2024-06-01', timestamp'1500-01-01 00:00:00'), " + + "(2, date'2024-06-01', timestamp'2024-06-01 12:00:00')") + } + withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC") { + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + val rows = df.selectExpr("id", "cast(ts as string)").collect().sortBy(_.getInt(0)) + assert( + rows(0).getString(1) == "1500-01-01 00:00:00", + s"$outputType/$datetimeMode/$int96Mode: got ${rows(0)}") + } + } + } + } + + /** + * Write one raw parquet file through parquet-mr's example writer: NO Spark writer metadata + * (`org.apache.spark.version` and friends) lands in the footer, the shape any non-Spark writer + * produces. Spark resolves such files' rebase policy from the session read modes + * (`DataSourceUtils.getRebaseSpec`'s `modeByConfig` fallback), so the native scan must too. + * Rows are (id, days-since-epoch date, micros-since-epoch UTC timestamp). + */ + private def writeNonSparkParquetFile( + dir: String, + rows: Seq[(Int, Option[Int], Option[Long])]): Unit = { + writeRawParquetFile( + dir, + """message m { + | required int32 id; + | optional int32 d (DATE); + | optional int64 ts (TIMESTAMP_MICROS); + |}""".stripMargin) { factory => + rows.map { case (id, d, ts) => + val group = factory.newGroup().append("id", id) + d.foreach(group.append("d", _)) + ts.foreach(group.append("ts", _)) + group + } + } + } + + /** + * Write one raw parquet file of the given parquet-mr `schema` (message type syntax) with the + * groups `rows` builds from a factory for that schema. Like [[writeNonSparkParquetFile]], no + * Spark writer metadata lands in the footer. + */ + private def writeRawParquetFile(dir: String, schema: String)( + rows: org.apache.parquet.example.data.simple.SimpleGroupFactory => Seq[ + org.apache.parquet.example.data.Group]): Unit = { + import org.apache.parquet.example.data.simple.SimpleGroupFactory + import org.apache.parquet.hadoop.example.{ExampleParquetWriter, GroupWriteSupport} + import org.apache.parquet.schema.MessageTypeParser + val messageType = MessageTypeParser.parseMessageType(schema) + val conf = new org.apache.hadoop.conf.Configuration() + GroupWriteSupport.setSchema(messageType, conf) + val writer = ExampleParquetWriter + .builder(new org.apache.hadoop.fs.Path(s"$dir/part-00000.parquet")) + .withConf(conf) + .build() + try { + rows(new SimpleGroupFactory(messageType)).foreach(writer.write) + } finally { + writer.close() + } + } + + /** + * The 12-byte INT96 encoding of midnight on the day `days` after 1970-01-01: 8 bytes of + * nanos-of-day then the 4-byte Julian Day Number (2440588 + days), both little-endian, the + * layout Spark's `ParquetRowConverter.binaryToSQLTimestamp` decodes. + */ + private def int96Midnight(days: Int): org.apache.parquet.io.api.Binary = { + val buf = java.nio.ByteBuffer.allocate(12).order(java.nio.ByteOrder.LITTLE_ENDIAN) + buf.putLong(0L).putInt(2440588 + days) + org.apache.parquet.io.api.Binary.fromConstantByteArray(buf.array()) + } + + /** Collect every message down the cause chain of `e`, newline-joined. */ + private def causeMessages(e: Throwable): String = + Iterator.iterate(e)(_.getCause).takeWhile(_ != null).map(_.getMessage).mkString("\n") + + /** Spark's `RebaseDateTime.lastSwitchJulianTs`: 1900-01-01T00:00:00Z in micros. */ + private val LastSwitchJulianMicros = -2208988800000000L + + test( + "non-Spark INT64 timestamps at or after 1900-01-01 read verbatim under EXCEPTION read " + + "modes") { + // Spark's EXCEPTION read mode refuses only timestamps before + // RebaseDateTime.lastSwitchJulianTs (1900-01-01T00:00:00Z, the last instant at which + // rebasing changes a value in any zone), converting MILLIS columns to micros first; a + // timestamp one microsecond before the epoch is well inside the accepted range. + withTempPath { dir => + val path = dir.getAbsolutePath + writeRawParquetFile( + path, + """message m { + | required int32 id; + | optional int64 ts_us (TIMESTAMP(MICROS,true)); + | optional int64 ts_ms (TIMESTAMP(MILLIS,true)); + |}""".stripMargin) { factory => + Seq( + factory.newGroup().append("id", 1).append("ts_us", -1L).append("ts_ms", -1L), + factory + .newGroup() + .append("id", 2) + .append("ts_us", LastSwitchJulianMicros) + .append("ts_ms", LastSwitchJulianMicros / 1000), + factory + .newGroup() + .append("id", 3) + .append("ts_us", 1717243200000000L) + .append("ts_ms", 1717243200000L), + factory.newGroup().append("id", 4)) + } + val table = "comet_nonspark_1900_" + java.util.UUID.randomUUID().toString.replace("-", "") + withTable(table) { + spark.sql( + s"CREATE TABLE $table (id INT, ts_us TIMESTAMP, ts_ms TIMESTAMP) USING PARQUET " + + s"LOCATION '$path'") + spark.sql(s"CONVERT TO DELTA $table NO STATISTICS") + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + "spark.sql.parquet.datetimeRebaseModeInRead" -> "EXCEPTION", + "spark.sql.parquet.int96RebaseModeInRead" -> "EXCEPTION") { + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + val rows = df + .selectExpr("id", "cast(ts_us as string)", "cast(ts_ms as string)") + .collect() + .sortBy(_.getInt(0)) + assert(rows(0).getString(1) == "1969-12-31 23:59:59.999999", s"got ${rows(0)}") + assert(rows(0).getString(2) == "1969-12-31 23:59:59.999", s"got ${rows(0)}") + assert(rows(1).getString(1) == "1900-01-01 00:00:00", s"got ${rows(1)}") + assert(rows(1).getString(2) == "1900-01-01 00:00:00", s"got ${rows(1)}") + assert(rows(2).getString(1) == "2024-06-01 12:00:00", s"got ${rows(2)}") + assert(rows(3).isNullAt(1) && rows(3).isNullAt(2), s"got ${rows(3)}") + } + } + } + } + + test("non-Spark INT64 timestamps before 1900-01-01 fail loudly under EXCEPTION read modes") { + // One millisecond before the cutoff, in a MILLIS column: Spark converts to micros before + // comparing against lastSwitchJulianTs and raises; the native scan must raise too. + withTempPath { dir => + val path = dir.getAbsolutePath + writeRawParquetFile( + path, + """message m { + | required int32 id; + | optional int64 ts_ms (TIMESTAMP(MILLIS,true)); + |}""".stripMargin) { factory => + Seq(factory.newGroup().append("id", 1).append("ts_ms", LastSwitchJulianMicros / 1000 - 1)) + } + val table = "comet_nonspark_1899_" + java.util.UUID.randomUUID().toString.replace("-", "") + withTable(table) { + spark.sql(s"CREATE TABLE $table (id INT, ts_ms TIMESTAMP) USING PARQUET LOCATION '$path'") + spark.sql(s"CONVERT TO DELTA $table NO STATISTICS") + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + "spark.sql.parquet.datetimeRebaseModeInRead" -> "EXCEPTION", + "spark.sql.parquet.int96RebaseModeInRead" -> "EXCEPTION") { + val e = intercept[Exception] { + spark.read.format("delta").load(path).collect() + } + val messages = causeMessages(e) + assert( + messages.contains("Native scan cannot rebase") && messages.contains("'ts_ms'"), + s"expected the native calendar-rebase error on ts_ms, got:\n$messages") + } + } + } + } + + /** A raw file with one INT64 MICROS timestamp (`ts`) and one INT96 timestamp (`ts96`). */ + private def writeInt64AndInt96File(dir: String, tsMicros: Long, int96Days: Int): Unit = { + writeRawParquetFile( + dir, + """message m { + | required int32 id; + | optional int64 ts (TIMESTAMP(MICROS,true)); + | optional int96 ts96; + |}""".stripMargin) { factory => + Seq( + factory + .newGroup() + .append("id", 1) + .append("ts", tsMicros) + .append("ts96", int96Midnight(int96Days)), + factory.newGroup().append("id", 2)) + } + } + + /** Proleptic 1500-01-01 as days / micros since the epoch. */ + private val AncientDays = -171664 + private val AncientMicros = AncientDays.toLong * 86400000000L + + test("non-Spark INT64 timestamps follow the datetime read mode when the INT96 mode differs") { + // Spark selects datetimeRebaseSpec for INT64 MICROS/MILLIS columns and int96RebaseSpec only + // for INT96 columns. Under datetime CORRECTED + int96 EXCEPTION an ancient INT64 value reads + // verbatim; it must not be refused just because the INT96 spec would refuse an ancient + // INT96 value (the INT96 column holds a modern one here). + withTempPath { dir => + val path = dir.getAbsolutePath + writeInt64AndInt96File(path, AncientMicros, int96Days = 19875) + val table = "comet_int64_vs_int96_" + java.util.UUID.randomUUID().toString.replace("-", "") + withTable(table) { + spark.sql( + s"CREATE TABLE $table (id INT, ts TIMESTAMP, ts96 TIMESTAMP) USING PARQUET " + + s"LOCATION '$path'") + spark.sql(s"CONVERT TO DELTA $table NO STATISTICS") + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + "spark.sql.parquet.datetimeRebaseModeInRead" -> "CORRECTED", + "spark.sql.parquet.int96RebaseModeInRead" -> "EXCEPTION") { + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + val rows = df + .selectExpr("id", "cast(ts as string)", "cast(ts96 as string)") + .collect() + .sortBy(_.getInt(0)) + assert(rows(0).getString(1) == "1500-01-01 00:00:00", s"got ${rows(0)}") + assert(rows(0).getString(2) == "2024-06-01 00:00:00", s"got ${rows(0)}") + assert(rows(1).isNullAt(1) && rows(1).isNullAt(2), s"got ${rows(1)}") + } + } + } + } + + test("non-Spark INT96 timestamps follow the INT96 read mode") { + withTempPath { dir => + val path = dir.getAbsolutePath + writeInt64AndInt96File(path, tsMicros = 0L, int96Days = AncientDays) + val table = "comet_int96_policy_" + java.util.UUID.randomUUID().toString.replace("-", "") + withTable(table) { + spark.sql( + s"CREATE TABLE $table (id INT, ts TIMESTAMP, ts96 TIMESTAMP) USING PARQUET " + + s"LOCATION '$path'") + spark.sql(s"CONVERT TO DELTA $table NO STATISTICS") + // datetime CORRECTED + int96 EXCEPTION: the ancient INT96 value is refused, naming the + // INT96 column (the INT64 column's epoch value is fine under either spec). + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + "spark.sql.parquet.datetimeRebaseModeInRead" -> "CORRECTED", + "spark.sql.parquet.int96RebaseModeInRead" -> "EXCEPTION") { + val e = intercept[Exception] { + spark.read.format("delta").load(path).collect() + } + val messages = causeMessages(e) + assert( + messages.contains("Native scan cannot rebase") && messages.contains("'ts96'"), + s"expected the native calendar-rebase error on ts96, got:\n$messages") + } + // Mirror image: datetime EXCEPTION + int96 CORRECTED reads the ancient INT96 value + // verbatim (Spark decodes the Julian Day Number directly, no calendar involved). + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + "spark.sql.parquet.datetimeRebaseModeInRead" -> "EXCEPTION", + "spark.sql.parquet.int96RebaseModeInRead" -> "CORRECTED") { + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + val rows = df + .selectExpr("id", "cast(ts as string)", "cast(ts96 as string)") + .collect() + .sortBy(_.getInt(0)) + assert(rows(0).getString(1) == "1970-01-01 00:00:00", s"got ${rows(0)}") + assert(rows(0).getString(2) == "1500-01-01 00:00:00", s"got ${rows(0)}") + } + } + } + } + + /** Proleptic 1800-01-01T00:00:00Z in days and micros: before Spark's 1900-01-01 cutoff. */ + private val Days1800 = -62091 + private val Micros1800 = Days1800.toLong * 86400000000L + + test("non-Spark tz-free INT64 timestamps read as TIMESTAMP follow the datetime read mode") { + // Spark's ParquetVectorUpdaterFactory keys on the requested type and checks only the unit + // of an INT64 timestamp annotation, so a TIMESTAMP(MICROS, isAdjustedToUTC=false) column + // read as TIMESTAMP goes through LongWithRebaseUpdater under datetimeRebaseModeInRead: + // EXCEPTION refuses the ancient value, CORRECTED reads it verbatim. + withTempPath { dir => + val path = dir.getAbsolutePath + writeRawParquetFile( + path, + """message m { + | required int32 id; + | optional int64 ts (TIMESTAMP(MICROS,false)); + |}""".stripMargin) { factory => + Seq( + factory.newGroup().append("id", 1).append("ts", Micros1800), + factory.newGroup().append("id", 2)) + } + val table = "comet_tzfree_as_ltz_" + java.util.UUID.randomUUID().toString.replace("-", "") + withTable(table) { + spark.sql(s"CREATE TABLE $table (id INT, ts TIMESTAMP) USING PARQUET LOCATION '$path'") + spark.sql(s"CONVERT TO DELTA $table NO STATISTICS") + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + "spark.sql.parquet.datetimeRebaseModeInRead" -> "EXCEPTION", + "spark.sql.parquet.int96RebaseModeInRead" -> "CORRECTED") { + val e = intercept[Exception] { + spark.read.format("delta").load(path).collect() + } + val messages = causeMessages(e) + assert( + messages.contains("Native scan cannot rebase") && messages.contains("'ts'"), + s"expected the native calendar-rebase error on ts, got:\n$messages") + // Spark's own reader refuses the same value under EXCEPTION. + withSQLConf(DeltaScanConf.COMET_DELTA_NATIVE_ENABLED.key -> "false") { + val sparkError = intercept[Exception] { + spark.read.format("delta").load(path).collect() + } + // SparkUpgradeException is private[spark], so match it by name. + val causes = Iterator.iterate(sparkError: Throwable)(_.getCause).takeWhile(_ != null) + assert( + causes.exists(_.getClass.getName == "org.apache.spark.SparkUpgradeException"), + s"expected Spark's own rebase error, got:\n${causeMessages(sparkError)}") + } + } + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + "spark.sql.parquet.datetimeRebaseModeInRead" -> "CORRECTED", + "spark.sql.parquet.int96RebaseModeInRead" -> "EXCEPTION") { + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + val rows = df.selectExpr("id", "cast(ts as string)").collect().sortBy(_.getInt(0)) + assert(rows(0).getString(1) == "1800-01-01 00:00:00", s"got ${rows(0)}") + assert(rows(1).isNullAt(1), s"got ${rows(1)}") + } + // LEGACY without a recorded writer zone needs the JVM's default zone, which the native + // scan cannot know: it refuses rather than guessing. + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + "spark.sql.parquet.datetimeRebaseModeInRead" -> "LEGACY", + "spark.sql.parquet.int96RebaseModeInRead" -> "CORRECTED") { + val e = intercept[Exception] { + spark.read.format("delta").load(path).collect() + } + val messages = causeMessages(e) + assert( + messages.contains("Native scan cannot rebase") && + messages.contains("timezone tables"), + s"expected the native writer-zone error on ts, got:\n$messages") + } + } + } + } + + test("INT96 and adjusted INT64 timestamps read as TIMESTAMP_NTZ are never rebased") { + // Spark 4.x reads INT96 as TIMESTAMP_NTZ through BinaryToSQLTimestampUpdater and adjusted + // INT64 through LongUpdater; neither consults a rebase mode, so ancient values read as + // stored even under EXCEPTION. Spark 3.x refuses these pairings up front (SPARK-36182). + assume(isSpark40Plus) + withTempPath { dir => + val path = dir.getAbsolutePath + writeRawParquetFile( + path, + """message m { + | required int32 id; + | optional int96 ts96; + | optional int64 ts64 (TIMESTAMP(MICROS,true)); + |}""".stripMargin) { factory => + Seq( + factory + .newGroup() + .append("id", 1) + .append("ts96", int96Midnight(Days1800)) + .append("ts64", Micros1800), + factory.newGroup().append("id", 2)) + } + val table = "comet_ltz_as_ntz_" + java.util.UUID.randomUUID().toString.replace("-", "") + withTable(table) { + spark.sql( + s"CREATE TABLE $table (id INT, ts96 TIMESTAMP_NTZ, ts64 TIMESTAMP_NTZ) USING PARQUET " + + s"LOCATION '$path'") + spark.sql(s"CONVERT TO DELTA $table NO STATISTICS") + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + "spark.sql.parquet.datetimeRebaseModeInRead" -> "EXCEPTION", + "spark.sql.parquet.int96RebaseModeInRead" -> "EXCEPTION") { + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + val rows = df + .selectExpr("id", "cast(ts96 as string)", "cast(ts64 as string)") + .collect() + .sortBy(_.getInt(0)) + assert(rows(0).getString(1) == "1800-01-01 00:00:00", s"got ${rows(0)}") + assert(rows(0).getString(2) == "1800-01-01 00:00:00", s"got ${rows(0)}") + assert(rows(1).isNullAt(1) && rows(1).isNullAt(2), s"got ${rows(1)}") + } + } + } + } + + test("ancient timestamps and dates read as BIGINT and INT are never rebase-checked") { + // Spark's ParquetVectorUpdaterFactory picks LongUpdater for INT64 read as BIGINT and + // IntegerUpdater for INT32 read as INT before any timestamp or date arm, so the stored + // integers come back unchanged even under EXCEPTION. NANOS columns read as BIGINT only with + // spark.sql.legacy.parquet.nanosAsLong. + val nanos1800 = Micros1800 * 1000L + withTempPath { dir => + val path = dir.getAbsolutePath + writeRawParquetFile( + path, + """message m { + | required int32 id; + | required int64 v (TIMESTAMP(NANOS,true)); + | optional int64 us (TIMESTAMP(MICROS,true)); + | optional int32 d (DATE); + |}""".stripMargin) { factory => + Seq( + factory + .newGroup() + .append("id", 1) + .append("v", nanos1800) + .append("us", Micros1800) + .append("d", Days1800), + factory.newGroup().append("id", 2).append("v", 1717243200000000000L)) + } + val table = "comet_temporal_as_int_" + java.util.UUID.randomUUID().toString.replace("-", "") + withSQLConf("spark.sql.legacy.parquet.nanosAsLong" -> "true") { + withTable(table) { + spark.sql( + s"CREATE TABLE $table (id INT, v BIGINT, us BIGINT, d INT) USING PARQUET " + + s"LOCATION '$path'") + spark.sql(s"CONVERT TO DELTA $table NO STATISTICS") + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + "spark.sql.parquet.datetimeRebaseModeInRead" -> "EXCEPTION", + "spark.sql.parquet.int96RebaseModeInRead" -> "EXCEPTION") { + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + val rows = df.collect().sortBy(_.getInt(0)) + assert(rows(0) == Row(1, nanos1800, Micros1800, Days1800), s"got ${rows(0)}") + assert(rows(1) == Row(2, 1717243200000000000L, null, null), s"got ${rows(1)}") + } + } + } + } + } + + test("a plain INT32 column read as DATE follows the datetime read mode") { + // Spark's ParquetVectorUpdaterFactory picks the INT32 updater from the requested type, so + // an INT32 column without a DATE annotation, declared DATE by the table, goes through + // IntegerWithRebaseUpdater under datetimeRebaseModeInRead like a DATE column does. + withTempPath { dir => + val path = dir.getAbsolutePath + writeRawParquetFile( + path, + """message m { + | required int32 id; + | optional int32 v; + |}""".stripMargin) { factory => + Seq( + factory.newGroup().append("id", 1).append("v", -171655), + factory.newGroup().append("id", 2).append("v", 19875), + factory.newGroup().append("id", 3)) + } + val table = "comet_int_as_date_" + java.util.UUID.randomUUID().toString.replace("-", "") + withTable(table) { + spark.sql(s"CREATE TABLE $table (id INT, v DATE) USING PARQUET LOCATION '$path'") + spark.sql(s"CONVERT TO DELTA $table NO STATISTICS") + Seq("LEGACY" -> "1500-01-01", "CORRECTED" -> "1500-01-10").foreach { + case (mode, expected) => + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + "spark.sql.parquet.datetimeRebaseModeInRead" -> mode) { + val df = spark.read.format("delta").load(path).selectExpr("id", "cast(v as string)") + checkDeltaNativeScanAnswer(df) + val rows = df.collect().sortBy(_.getInt(0)).toSeq + assert( + rows == Seq(Row(1, expected), Row(2, "2024-06-01"), Row(3, null)), + s"$mode: got $rows") + } + } + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + "spark.sql.parquet.datetimeRebaseModeInRead" -> "EXCEPTION") { + val e = intercept[Exception] { + spark.read.format("delta").load(path).collect() + } + val causes = Iterator.iterate(e: Throwable)(_.getCause).takeWhile(_ != null).toSeq + assert( + causes.exists(c => + c.getClass.getName == "org.apache.comet.CometNativeException" && + c.getMessage.contains("Native scan cannot rebase") && + c.getMessage.contains("'v'")), + s"expected the native calendar-rebase error on v, got:\n${causeMessages(e)}") + withSQLConf(DeltaScanConf.COMET_DELTA_NATIVE_ENABLED.key -> "false") { + val sparkError = intercept[Exception] { + spark.read.format("delta").load(path).collect() + } + val sparkCauses = + Iterator.iterate(sparkError: Throwable)(_.getCause).takeWhile(_ != null) + assert( + sparkCauses.exists(_.getClass.getName == "org.apache.spark.SparkUpgradeException"), + s"expected Spark's own rebase error, got:\n${causeMessages(sparkError)}") + } + } + } + } + } + + private val NestedRawSchema = """message m { + | required int32 id; + | optional group s { + | optional int32 d (DATE); + | optional int64 ts (TIMESTAMP(MICROS,true)); + | } + | optional group l (LIST) { + | repeated group list { + | optional int32 element (DATE); + | } + | } + |}""".stripMargin + + private def createNestedRawTable(table: String, path: String): Unit = { + spark.sql( + s"CREATE TABLE $table (id INT, s STRUCT, l ARRAY) " + + s"USING PARQUET LOCATION '$path'") + spark.sql(s"CONVERT TO DELTA $table NO STATISTICS") + } + + test( + "metadata-free nested columns with modern and null datetime leaves stay native under " + + "EXCEPTION read modes") { + // EXCEPTION only refuses values that actually are ancient; a STRUCT + // and an ARRAY holding modern and null leaves must read natively, not be rejected + // up front for being nested. + withTempPath { dir => + val path = dir.getAbsolutePath + writeRawParquetFile(path, NestedRawSchema) { factory => + val g1 = factory.newGroup().append("id", 1) + g1.addGroup("s").append("d", 19875).append("ts", 1717243200000000L) + val l1 = g1.addGroup("l") + l1.addGroup("list").append("element", 19875) + l1.addGroup("list") + val g2 = factory.newGroup().append("id", 2) + g2.addGroup("s") + g2.addGroup("l") + val g3 = factory.newGroup().append("id", 3) + Seq(g1, g2, g3) + } + val table = "comet_nested_modern_" + java.util.UUID.randomUUID().toString.replace("-", "") + withTable(table) { + createNestedRawTable(table, path) + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + "spark.sql.parquet.datetimeRebaseModeInRead" -> "EXCEPTION", + "spark.sql.parquet.int96RebaseModeInRead" -> "EXCEPTION") { + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + val rows = df + .selectExpr("id", "cast(s.d as string)", "cast(s.ts as string)", "cast(l as string)") + .collect() + .sortBy(_.getInt(0)) + assert(rows(0).getString(1) == "2024-06-01", s"got ${rows(0)}") + assert(rows(0).getString(2) == "2024-06-01 12:00:00", s"got ${rows(0)}") + assert(rows(0).getString(3) == "[2024-06-01, null]", s"got ${rows(0)}") + assert(rows(1).isNullAt(1) && rows(1).isNullAt(2), s"got ${rows(1)}") + assert(rows(1).getString(3) == "[]", s"got ${rows(1)}") + assert(rows(2).isNullAt(1) && rows(2).isNullAt(3), s"got ${rows(2)}") + } + } + } + } + + test( + "metadata-free nested columns with an ancient date leaf fail loudly under EXCEPTION read " + + "modes") { + withTempPath { dir => + val path = dir.getAbsolutePath + writeRawParquetFile(path, NestedRawSchema) { factory => + val g1 = factory.newGroup().append("id", 1) + g1.addGroup("s").append("d", -171655) + Seq(g1) + } + val table = "comet_nested_ancient_" + java.util.UUID.randomUUID().toString.replace("-", "") + withTable(table) { + createNestedRawTable(table, path) + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + "spark.sql.parquet.datetimeRebaseModeInRead" -> "EXCEPTION", + "spark.sql.parquet.int96RebaseModeInRead" -> "EXCEPTION") { + val e = intercept[Exception] { + spark.read.format("delta").load(path).collect() + } + val messages = causeMessages(e) + assert( + messages.contains("Native scan cannot rebase") && messages.contains("'s'"), + s"expected the native calendar-rebase error on s, got:\n$messages") + } + } + } + } + + test( + "only the requested nested leaves are rebase-checked: an unrequested ancient s.ts does not " + + "block select s.d under EXCEPTION read modes") { + // A metadata-free file with s.d = 2024-06-01 next to s.ts = 1500-01-01. Spark's requested + // schema for `select s.d` is STRUCT, so Spark never decodes s.ts and reads the modern + // date fine; the native scan must not refuse the row for a leaf the schema adapter's struct + // narrowing drops. Requesting the ancient leaf itself still fails loudly. + withTempPath { dir => + val path = dir.getAbsolutePath + writeRawParquetFile(path, NestedRawSchema) { factory => + val g1 = factory.newGroup().append("id", 1) + g1.addGroup("s").append("d", 19875).append("ts", AncientMicros) + Seq(g1) + } + val table = + "comet_nested_requested_" + java.util.UUID.randomUUID().toString.replace("-", "") + withTable(table) { + createNestedRawTable(table, path) + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + "spark.sql.parquet.datetimeRebaseModeInRead" -> "EXCEPTION", + "spark.sql.parquet.int96RebaseModeInRead" -> "EXCEPTION") { + val df = spark.read.format("delta").load(path).selectExpr("id", "cast(s.d as string)") + checkDeltaNativeScanAnswer(df) + val rows = df.collect() + assert(rows.length == 1 && rows(0).getString(1) == "2024-06-01", s"got ${rows.toSeq}") + + for (projection <- Seq("s", "s.ts")) { + val e = intercept[Exception] { + spark.read.format("delta").load(path).selectExpr(projection).collect() + } + val messages = causeMessages(e) + assert( + messages.contains("Native scan cannot rebase") && messages.contains("'s'"), + s"expected the native calendar-rebase error on s for `select $projection`, " + + s"got:\n$messages") + } + } + } + } + } + + test("legacy-rebase ancient datetime values inside nested columns match Spark's own read") { + // Spark rebases dates and timestamps at every nesting depth; a LEGACY (UTC) file with + // ancient leaves inside a struct, an array, a map and an array of structs must read + // natively with exactly Spark's rebased values, nulls and offsets preserved. + withTempPath { dir => + val path = dir.getAbsolutePath + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + "spark.sql.parquet.datetimeRebaseModeInWrite" -> "LEGACY", + "spark.sql.parquet.int96RebaseModeInWrite" -> "LEGACY") { + spark + .sql("select * from values " + + "(1, named_struct('d', date'1500-01-01', 'ts', timestamp'1500-01-01 12:34:56'), " + + "array(date'1500-01-01', null, date'2024-06-01'), " + + "map(1, date'1582-10-04', 2, cast(null as date)), " + + "array(named_struct('d', date'0001-01-01'), named_struct('d', cast(null as date)))), " + + "(2, named_struct('d', cast(null as date), 'ts', cast(null as timestamp)), " + + "array(), map(), array(cast(null as struct))), " + + "(3, cast(null as struct), cast(null as array), " + + "cast(null as map), cast(null as array>)) " + + "as t(id, s, l, m, ls)") + .write + .format("delta") + .save(path) + } + withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC") { + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + val rows = df + .selectExpr( + "id", + "cast(s.d as string)", + "cast(s.ts as string)", + "cast(l as string)", + "cast(m as string)", + "cast(ls as string)") + .collect() + .sortBy(_.getInt(0)) + assert(rows(0).getString(1) == "1500-01-01", s"got ${rows(0)}") + assert(rows(0).getString(2) == "1500-01-01 12:34:56", s"got ${rows(0)}") + assert(rows(0).getString(3) == "[1500-01-01, null, 2024-06-01]", s"got ${rows(0)}") + assert(rows(0).getString(4) == "{1 -> 1582-10-04, 2 -> null}", s"got ${rows(0)}") + assert(rows(0).getString(5) == "[{0001-01-01}, {null}]", s"got ${rows(0)}") + assert(rows(1).isNullAt(1) && rows(1).isNullAt(2), s"got ${rows(1)}") + assert(rows(1).getString(3) == "[]" && rows(1).getString(4) == "{}", s"got ${rows(1)}") + assert(rows(1).getString(5) == "[null]", s"got ${rows(1)}") + assert((1 to 5).forall(rows(2).isNullAt), s"got ${rows(2)}") + } + } + } + + /** Register `path`'s raw parquet files as an external table and CONVERT it to Delta. */ + private def convertRawParquetToDelta(path: String, table: String): Unit = { + spark.sql( + s"CREATE TABLE $table (id INT, d DATE, ts TIMESTAMP) USING PARQUET LOCATION '$path'") + spark.sql(s"CONVERT TO DELTA $table NO STATISTICS") + } + + test("non-Spark parquet files read ancient values verbatim under CORRECTED read modes") { + // A converted table over a file with no Spark writer metadata: getRebaseSpec resolves the + // policy from the session read modes (Spark 4.0 defaults both to CORRECTED), so a + // proleptic 1500-01-01 (day -171664) and a timestamp one microsecond before the epoch + // must read natively exactly as stored. + withTempPath { dir => + val path = dir.getAbsolutePath + writeNonSparkParquetFile( + path, + Seq( + (1, Some(-171664), Some(-1L)), + (2, Some(19875), Some(1717243200000000L)), + (3, None, None))) + val table = + "comet_nonspark_corrected_" + java.util.UUID.randomUUID().toString.replace("-", "") + withTable(table) { + convertRawParquetToDelta(path, table) + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + "spark.sql.parquet.datetimeRebaseModeInRead" -> "CORRECTED", + "spark.sql.parquet.int96RebaseModeInRead" -> "CORRECTED") { + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + val rows = df + .selectExpr("id", "cast(d as string)", "cast(ts as string)") + .collect() + .sortBy(_.getInt(0)) + assert(rows(0).getString(1) == "1500-01-01", s"got ${rows(0)}") + assert(rows(0).getString(2) == "1969-12-31 23:59:59.999999", s"got ${rows(0)}") + assert(rows(1).getString(1) == "2024-06-01", s"got ${rows(1)}") + assert(rows(2).isNullAt(1) && rows(2).isNullAt(2), s"got ${rows(2)}") + } + } + } + } + + test("non-Spark parquet files rebase ancient dates under LEGACY read modes") { + // LEGACY read modes on a file without writer metadata: the stored day count is hybrid + // Julian + Gregorian, so Julian 1500-01-01 (stored as -171655) must rebase to proleptic + // 1500-01-01, matching Spark's own LEGACY read (the day rebase is timezone-free). + // Timestamps stay modern: rebasing ancient ones needs the writer zone, which this file + // does not record. + withTempPath { dir => + val path = dir.getAbsolutePath + writeNonSparkParquetFile( + path, + Seq((1, Some(-171655), Some(0L)), (2, Some(19875), Some(1717243200000000L)))) + val table = "comet_nonspark_legacy_" + java.util.UUID.randomUUID().toString.replace("-", "") + withTable(table) { + convertRawParquetToDelta(path, table) + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + "spark.sql.parquet.datetimeRebaseModeInRead" -> "LEGACY", + "spark.sql.parquet.int96RebaseModeInRead" -> "LEGACY") { + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + val rows = df + .selectExpr("id", "cast(d as string)") + .collect() + .sortBy(_.getInt(0)) + assert(rows(0).getString(1) == "1500-01-01", s"got ${rows(0)}") + assert(rows(1).getString(1) == "2024-06-01", s"got ${rows(1)}") + } + } + } + } + + test("non-Spark parquet files with ancient values fail loudly under EXCEPTION read modes") { + // EXCEPTION read modes (Spark 3.x's default) refuse ancient values whose calendar the + // file does not declare; the native scan must refuse them too rather than return + // silently shifted values. + withTempPath { dir => + val path = dir.getAbsolutePath + writeNonSparkParquetFile(path, Seq((1, Some(-171655), Some(0L)))) + val table = + "comet_nonspark_exception_" + java.util.UUID.randomUUID().toString.replace("-", "") + withTable(table) { + convertRawParquetToDelta(path, table) + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + "spark.sql.parquet.datetimeRebaseModeInRead" -> "EXCEPTION", + "spark.sql.parquet.int96RebaseModeInRead" -> "EXCEPTION") { + val e = intercept[Exception] { + spark.read.format("delta").load(path).collect() + } + val messages = Iterator + .iterate(e: Throwable)(_.getCause) + .takeWhile(_ != null) + .map(_.getMessage) + .mkString("\n") + assert( + messages.toLowerCase(java.util.Locale.ROOT).contains("rebase"), + s"expected a calendar-rebase error, got:\n$messages") + } + } + } + } + + test( + "a struct with a date column stays native when the file carries only the legacy INT96 " + + "flag and corrected dates") { + // legacyINT96 alone puts the file's INT96 timestamp column under the LEGACY policy, but + // its DATE policy is CORRECTED -- so a STRUCT column has nothing to rebase and + // must pass through natively, unwrapped, instead of being handled just because the + // timestamp policy needs handling elsewhere in the file. + withTempPath { dir => + val path = dir.getAbsolutePath + withSQLConf( + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC", + "spark.sql.parquet.outputTimestampType" -> "INT96", + "spark.sql.parquet.datetimeRebaseModeInWrite" -> "CORRECTED", + "spark.sql.parquet.int96RebaseModeInWrite" -> "LEGACY") { + spark + .sql( + "select * from values " + + "(1, named_struct('d', date'2020-06-01'), timestamp'2021-01-01 00:00:00'), " + + "(2, named_struct('d', cast(null as date)), timestamp'2022-01-01 12:34:56') " + + "as t(id, s, ts)") + .write + .format("delta") + .save(path) + } + withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC") { + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + val rows = df.selectExpr("id", "cast(s.d as string)").collect().sortBy(_.getInt(0)) + assert(rows(0).getString(1) == "2020-06-01", s"got ${rows(0)}") + assert(rows(1).isNullAt(1), s"got ${rows(1)}") + } + } + } +} diff --git a/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaS3Suite.scala b/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaS3Suite.scala new file mode 100644 index 00000000000..4f758e2430e --- /dev/null +++ b/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaS3Suite.scala @@ -0,0 +1,307 @@ +/* + * 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.contrib.delta + +import java.util.Locale + +import scala.util.{Failure, Success, Try} + +import org.testcontainers.DockerClientFactory + +import org.apache.hadoop.fs.Path +import org.apache.spark.internal.Logging +import org.apache.spark.sql.delta.DeltaLog +import org.apache.spark.sql.delta.actions.DeletionVectorDescriptor + +import org.apache.comet.CometS3TestBase + +/** + * MinIO-backed integration coverage for multi-bucket Delta shapes: a real two-bucket shallow + * clone, which a single-bucket `withTempPath` table can never produce, because it is + * `DeltaTable`'s CLONE machinery -- not test fixturing -- that leaves some `AddFile` entries + * pointing at the source table's absolute location while new files land under the clone's own + * root. + * + * Manual/opt-in, same as [[org.apache.comet.parquet.ParquetReadFromS3Suite]] in the spark module + * -- but gated differently out of necessity. That suite is invisible to every PR workflow simply + * because `.github/workflows/pr_build_linux.yml` / `pr_build_macos.yml` enumerate test classes by + * name and never name it (`dev/ci/check-suites.py` exempts it via `ignore_list` instead of + * requiring it be listed). The contrib module has no such allowlist: `delta_contrib_test.yml` + * runs `mvn ... test -pl contrib/delta-spark`, which discovers and runs every suite on the + * module's test classpath, and `check-suites.py` does not enforce anything under `contrib/` at + * all (see its `path.parts[0] == "contrib"` skip), so there is no file to omit this suite from. + * Every test therefore starts with `assume(dockerAvailable, ...)`: when no Docker daemon is + * reachable, ScalaTest reports the test CANCELED rather than failed or run, which + * `scalatest-maven-plugin` does not treat as a build failure -- the practical equivalent of + * `ParquetReadFromS3Suite`'s blanket omission, reached by a runtime check instead of never being + * named. `beforeAll` mirrors this: it probes Docker BEFORE calling `CometS3TestBase#beforeAll`, + * because that trait's `sparkConf` dereferences `minioContainer` unconditionally, and starting + * the Spark session (let alone a container) is exactly what a Docker-less run must not do. + * + * That fail-soft default makes a zero-coverage run look green, so the CI job that exists only to + * run this suite sets `COMET_DELTA_S3_REQUIRED=1`, which turns a missing Docker daemon or a + * failed MinIO start into a thrown `beforeAll`: the suite aborts and `scalatest-maven-plugin` + * fails. + */ +class CometDeltaS3Suite extends CometDeltaTestBase with CometS3TestBase with Logging { + + override protected val testBucketName = "comet-delta-a" + + /** + * The clone's destination bucket: distinct from [[testBucketName]] on purpose -- these tests + * exist to put a table's data (or its deletion vectors) across two object-store authorities. + */ + private val cloneBucketName = "comet-delta-b" + + /** + * A bucket touched by no other test in this suite: the native S3 object-store cache + * (`object_store_cache` in parquet_support.rs) is process-wide and keyed per bucket, so reusing + * [[testBucketName]] for the `${...}` forwarding test below risks silently passing against a + * store handle another test already warmed with plain credentials, rather than actually forcing + * a fresh credential derivation through the substituted `${...}` value. + */ + private val reviewRefBucketName = "comet-delta-review-ref" + + private var dockerAvailable = false + + override def beforeAll(): Unit = { + val required = CometDeltaS3Suite.s3Required(sys.env.get(CometDeltaS3Suite.S3_REQUIRED_ENV)) + dockerAvailable = DockerClientFactory.instance().isDockerAvailable + if (!dockerAvailable && required) { + throw new IllegalStateException( + CometDeltaS3Suite.requiredFailureMessage("no Docker daemon is reachable")) + } + if (dockerAvailable) { + // Fail soft unless COMET_DELTA_S3_REQUIRED arms the hard failure: this suite runs + // unconditionally in CI (no allowlist to omit it from, see the class doc above), and + // testcontainers networking inside a CI job container is unverified -- MinIO is a sibling + // container there, so `getS3URL` may resolve to an address that is wrong from inside the + // job container. If startup or bucket creation blows up, log the resolved URL (the signal + // needed to diagnose a first bad CI run), flip `dockerAvailable` back off so every test + // cancels via `assume` instead of aborting the whole suite, and best-effort stop whatever + // container did come up. + Try { + super.beforeAll() // CometS3TestBase starts MinIO, then CometTestBase starts the session. + createBucketIfNotExists(cloneBucketName) + createBucketIfNotExists(reviewRefBucketName) + } match { + case Success(_) => + logInfo(s"CometDeltaS3Suite: MinIO reachable at ${minioContainer.getS3URL}") + case Failure(e) => + val resolvedUrl = Try(minioContainer.getS3URL).getOrElse("") + val cause = s"MinIO setup failed (resolved S3 URL: $resolvedUrl)" + dockerAvailable = false + // Tear down here, synchronously: super.beforeAll() may have partially succeeded + // (e.g. the Spark session started but createBucketIfNotExists(cloneBucketName) + // failed), and this suite's own afterAll() below is gated on `dockerAvailable`, + // which is now false -- the framework-invoked afterAll() will no-op and never get a + // chance to stop anything. super.afterAll() stops both the Spark session + // (CometTestBase#afterAll, tolerates a session that never started) and MinIO + // (CometS3TestBase#afterAll, tolerates a container that never started), so this is + // safe to call unconditionally here regardless of how far beforeAll got. + Try(super.afterAll()) + if (required) { + throw new IllegalStateException(CometDeltaS3Suite.requiredFailureMessage(cause), e) + } + logWarning(s"CometDeltaS3Suite: $cause; skipping all tests in this suite", e) + } + } + } + + override def afterAll(): Unit = { + if (dockerAvailable) { + super.afterAll() + } + } + + // CometTestBase#afterEach unconditionally touches `spark` (cache-clearing, open-stream + // assertions); with no session ever created in a Docker-less run, that NPEs and aborts the + // whole suite -- turning a clean per-test cancellation into a module-wide build failure. + override def afterEach(): Unit = { + if (dockerAvailable) { + super.afterEach() + } + } + + private def tablePath(bucket: String, relPath: String): String = s"s3a://$bucket/$relPath" + + test("shallow clone across buckets + append declines with the multi-store reason") { + assume(dockerAvailable, "Docker is not available; skipping MinIO-backed Delta test") + + val sourcePath = tablePath(testBucketName, "clone-append/source") + val clonePath = tablePath(cloneBucketName, "clone-append/clone") + + spark.range(0, 100).selectExpr("id", "id * 2 as v").write.format("delta").save(sourcePath) + spark.sql(s"CREATE TABLE delta.`$clonePath` SHALLOW CLONE delta.`$sourcePath`") + // The clone's own transaction log still references the SOURCE's physical files (bucket A) + // for every row carried over by the clone. This append writes NEW physical files under the + // clone's own root (bucket B): the clone's data files now span two object-store + // authorities -- exactly the shape the multi-store decline gate exists for, since the shared + // native scan builder resolves the whole scan's ObjectStoreUrl from the first selected file + // only. + spark + .range(100, 150) + .selectExpr("id", "id * 2 as v") + .write + .format("delta") + .mode("append") + .save(clonePath) + + val df = spark.read.format("delta").load(clonePath) + checkSparkAnswerAndFallbackReason( + df, + "Native Delta scan does not support data files spanning multiple object stores") + } + + test( + "clone across buckets + DELETE on the clone reads correct rows natively " + + "(cold cross-bucket deletion-vector store)") { + assume(dockerAvailable, "Docker is not available; skipping MinIO-backed Delta test") + + val sourcePath = tablePath(testBucketName, "clone-delete/source") + val clonePath = tablePath(cloneBucketName, "clone-delete/clone") + + spark.range(0, 1000).selectExpr("id", "id * 2 as v").write.format("delta").save(sourcePath) + spark.sql( + s"ALTER TABLE delta.`$sourcePath` SET TBLPROPERTIES ('delta.enableDeletionVectors' = 'true')") + spark.sql(s"CREATE TABLE delta.`$clonePath` SHALLOW CLONE delta.`$sourcePath`") + // DELETE against a deletion-vector table does not rewrite the target file; it attaches a + // deletion-vector sidecar to the existing `AddFile` action instead. The sidecar is written + // under the CLONE's own root (bucket B), while the `AddFile` it decorates still points at + // the SOURCE's absolute, un-copied physical file (bucket A) -- shallow clone never + // relocates data it did not modify. That is the cold cross-bucket deletion-vector-store bug + // shape: attaching the access plan nested a `Handle::block_on` call that built the + // (previously untouched, so cold) bucket-B object store for the sidecar from inside an + // already-running Tokio runtime, which panics. + // + // It is also, deliberately, NOT the shape the multi-store decline gate catches: that gate + // inspects only DATA-file authorities (`scanHelper.selectedPartitions...map(_.getPath)`), + // and every data file this scan selects is still on bucket A -- only the deletion vector's + // own authority is bucket B. That distinction is worth stating here because it is + // the one property that makes this test exercise the cold cross-bucket deletion-vector-store + // path instead of re-proving the multi-store decline gate. + spark.sql(s"DELETE FROM delta.`$clonePath` WHERE id % 2 = 0") + + // Assert the cross-bucket shape STRUCTURALLY, not just end-to-end via the read below: if + // Delta's shallow-clone or DELETE-on-a-DV-table semantics ever change (DELETE starts + // rewriting the file instead of writing a DV, or the DV sidecar starts landing next to the + // data it decorates instead of under the clone's own root), the test must fail loudly right + // here -- otherwise it would silently degrade into a same-bucket read that never exercises + // the cold cross-bucket deletion-vector-store code path at all, while + // `checkDeltaNativeScanAnswer` below would still pass. + val log = DeltaLog.forTable(spark, clonePath) + val cloneTableRootPath = new Path(clonePath) + val files = log.update().allFiles.collect() + + // At least one data file must still resolve into the SOURCE bucket: shallow clone never + // copies files it did not modify. + val dataAuthorities = files.map(_.absolutePath(log).toUri.getHost).distinct + assert( + dataAuthorities.contains(testBucketName), + "expected at least one data file to still resolve into the SOURCE bucket " + + s"($testBucketName, carried over unmodified by the shallow clone); resolved data-file " + + s"authorities: ${dataAuthorities.mkString(", ")}") + + // At least one deletion-vector descriptor must resolve into the CLONE's own bucket. + // Resolution mirrors DeltaScanSupport.selectedDvDescriptors (copyWithAbsolutePath against + // the table root) followed by CometDeltaNativeScan.storeUris's own absolutePath call -- + // the exact path production code takes from AddFile to an object-store authority. Inline + // or canonically-empty descriptors are excluded first (`cardinality == 0` is the + // EMPTY-descriptor characterization: no rows deleted, so no on-disk sidecar exists): + // DeletionVectorDescriptor#absolutePath's isOnDisk precondition + // throws for inline ones, and neither carries a resolvable external authority. + val dvAuthorities = files + .flatMap(f => Option(f.deletionVector)) + .filter(dv => + dv.storageType != DeletionVectorDescriptor.INLINE_DV_MARKER && dv.cardinality > 0) + .map( + _.copyWithAbsolutePath(cloneTableRootPath).absolutePath(cloneTableRootPath).toUri.getHost) + .distinct + assert( + dvAuthorities.contains(cloneBucketName), + "expected at least one deletion-vector sidecar to resolve into the CLONE's own bucket " + + s"($cloneBucketName); resolved deletion-vector authorities: " + + s"${dvAuthorities.mkString(", ")}") + + val df = spark.read.format("delta").load(clonePath) + checkDeltaNativeScanAnswer(df) + assert(df.count() == 500) + } + + test( + s"S3 credentials configured via a Hadoop $${...} variable reference (fs.s3a.access.key = " + + s"$${review.access}, fs.s3a.secret.key = $${review.secret}) claim natively and read " + + "correct rows against a real MinIO bucket") { + assume(dockerAvailable, "Docker is not available; skipping MinIO-backed Delta test") + + // Mutate the session's shared hadoopConfiguration directly (mirroring the set-then-restore + // shape CometDeltaNativeScanSuite's viewfs gate test uses for the same reason: these are + // plain Hadoop Configuration entries, not SQLConf, so withSQLConf cannot round-trip them). + // Aliasing the real MinIO credentials behind review.access/review.secret and pointing + // fs.s3a.access.key/fs.s3a.secret.key at them via ${...} reproduces exactly the shape + // PART 1 fixed: Configuration#get expands the reference to the real credential, and + // NativeConfig.extractObjectStoreOptions must forward that EXPANDED value, not the literal + // "${review.access}" string, or the native S3 client would authenticate with garbage. + val hadoopConf = spark.sparkContext.hadoopConfiguration + val priorAccessKey = Option(hadoopConf.get("fs.s3a.access.key")) + val priorSecretKey = Option(hadoopConf.get("fs.s3a.secret.key")) + hadoopConf.set("review.access", userName) + hadoopConf.set("review.secret", password) + hadoopConf.set("fs.s3a.access.key", s"$${review.access}") + hadoopConf.set("fs.s3a.secret.key", s"$${review.secret}") + try { + val path = tablePath(reviewRefBucketName, "review-ref-table") + spark.range(0, 200).selectExpr("id", "id * 2 as v").write.format("delta").save(path) + + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + assert(df.count() == 200) + } finally { + hadoopConf.unset("review.access") + hadoopConf.unset("review.secret") + priorAccessKey match { + case Some(v) => hadoopConf.set("fs.s3a.access.key", v) + case None => hadoopConf.unset("fs.s3a.access.key") + } + priorSecretKey match { + case Some(v) => hadoopConf.set("fs.s3a.secret.key", v) + case None => hadoopConf.unset("fs.s3a.secret.key") + } + } + } +} + +object CometDeltaS3Suite { + + /** Environment variable that turns a Docker-less or MinIO-less run into a suite failure. */ + val S3_REQUIRED_ENV = "COMET_DELTA_S3_REQUIRED" + + /** + * Whether the run must fail rather than cancel when MinIO is unavailable. Strict on purpose: + * only `1` and `true` (trimmed, case-insensitive) arm it; anything else, including `yes` and + * `0`, keeps the fail-soft default so a typo cannot arm or disarm the switch unnoticed. + */ + private[delta] def s3Required(value: Option[String]): Boolean = + value.map(_.trim.toLowerCase(Locale.ROOT)).exists(v => v == "1" || v == "true") + + /** Names the env var and the cause so a red CI run reads directly from the failure line. */ + private[delta] def requiredFailureMessage(cause: String): String = + s"$S3_REQUIRED_ENV is set but $cause; failing CometDeltaS3Suite instead of cancelling it" +} diff --git a/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaTestBase.scala b/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaTestBase.scala new file mode 100644 index 00000000000..50fbcc23fe2 --- /dev/null +++ b/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaTestBase.scala @@ -0,0 +1,57 @@ +/* + * 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.contrib.delta + +import org.apache.spark.SparkConf +import org.apache.spark.sql.{CometTestBase, DataFrame} +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + +/** + * Base for Delta contrib suites: CometTestBase plus the Delta Lake session extension and catalog. + */ +abstract class CometDeltaTestBase extends CometTestBase with AdaptiveSparkPlanHelper { + + override protected def sparkConf: SparkConf = { + val conf = super.sparkConf + conf.set("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension") + conf.set("spark.sql.catalog.spark_catalog", "org.apache.spark.sql.delta.catalog.DeltaCatalog") + conf.set(DeltaScanConf.COMET_DELTA_NATIVE_ENABLED.key, "true") + conf + } + + /** Collect nodes of the given simple class name anywhere in the (AQE-stripped) plan. */ + protected def collectByName(plan: SparkPlan, simpleName: String): Seq[SparkPlan] = + collectWithSubqueries(stripAQEPlan(plan)) { + case op if op.getClass.getSimpleName == simpleName => op + } + + protected def deltaNativeScans(df: DataFrame): Seq[SparkPlan] = + collectByName(df.queryExecution.executedPlan, "CometDeltaNativeScanExec") + + /** Assert the query ran through the native Delta scan AND matches the comet-off answer. */ + protected def checkDeltaNativeScanAnswer(df: DataFrame): Unit = { + checkSparkAnswer(df) + // Re-materialize the plan after execution so AQE has finalized stages. + assert( + deltaNativeScans(df).nonEmpty, + s"Expected CometDeltaNativeScanExec in plan:\n${df.queryExecution.executedPlan}") + } +} diff --git a/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/DeltaScanContribSuite.scala b/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/DeltaScanContribSuite.scala new file mode 100644 index 00000000000..6d45bf7a568 --- /dev/null +++ b/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/DeltaScanContribSuite.scala @@ -0,0 +1,2905 @@ +/* + * 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.contrib.delta + +import java.io.{File, IOException} +import java.net.URI +import java.nio.file.Files +import java.util.{Locale, UUID} +import java.util.concurrent.atomic.AtomicBoolean + +import scala.jdk.CollectionConverters._ + +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.{Path, RawLocalFileSystem} +import org.apache.hadoop.fs.s3a.{S3AFileSystem, S3AUtils} +import org.apache.hadoop.security.alias.CredentialProviderFactory +import org.apache.spark.sql.comet.CometDeltaNativeScanExec +import org.apache.spark.sql.delta.actions.DeletionVectorDescriptor + +import org.apache.comet.{CometConf, ExtendedExplainInfo} +import org.apache.comet.objectstore.NativeConfig +import org.apache.comet.rules.CometScanRule + +/** + * Guards the contrib claim path: the contrib is never active when Comet exec or Comet scan is + * disabled, and the claim hook runs before core's metadata-column guard. + */ +class DeltaScanContribSuite extends CometDeltaTestBase { + + test("contrib is inert when comet exec is disabled") { + // The COMET_EXEC_ENABLED gate lives in DeltaScanContrib.tryTransformV1; this pins it there. + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 100).write.format("delta").save(path) + + withSQLConf(CometConf.COMET_EXEC_ENABLED.key -> "false") { + val df = spark.read.format("delta").load(path) + checkSparkAnswer(df) + assert(deltaNativeScans(df).isEmpty) + } + } + } + + test("contrib is inert when comet native scan is disabled") { + // COMET_NATIVE_SCAN_ENABLED is checked in CometScanRule.transformScan before any V1 + // handling, so it short-circuits the CometScanContrib hook too. + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 100).write.format("delta").save(path) + + withSQLConf(CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "false") { + val df = spark.read.format("delta").load(path) + checkSparkAnswer(df) + assert(deltaNativeScans(df).isEmpty) + } + } + } + + test("claim runs before core's metadata-column guard") { + // A DV read's plan carries generated metadata columns that core's generic V1 guard + // would decline; the scan still goes native because CometScanContrib.tryTransformV1 + // is consulted first (CometScanRule.transformV1Scan hook order). + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 1000).selectExpr("id", "id * 2 as v").write.format("delta").save(path) + spark.sql( + s"ALTER TABLE delta.`$path` SET TBLPROPERTIES ('delta.enableDeletionVectors' = 'true')") + spark.sql(s"DELETE FROM delta.`$path` WHERE id % 2 = 0") + + val df = spark.read.format("delta").load(path) + checkSparkAnswer(df) + assert(deltaNativeScans(df).nonEmpty) + } + } + + test("declined scan carries the contrib's fallback reason, not core's generic one") { + // Disabling Spark's vectorized Parquet reader is a scan the contrib recognizes + // (DeltaScanSupport.isDeltaScan) but explicitly declines (DeltaScanSupport.declineReason, + // mirroring core's own vectorized-reader gate). Per the CometScanContrib ownership + // contract the contrib still claims it (tagging its own fallback reason), so core's + // generic V1 gate -- and its "Unsupported file format" message -- never runs on it. + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 100).write.format("delta").save(path) + + withSQLConf( + "spark.sql.parquet.enableVectorizedReader" -> "false", + // CometTestBase flips this to "true" so the rest of the suite can exercise the + // vectorized-off path against Comet's native scan; put it back to its real default so + // this gate actually declines. + CometConf.COMET_SCAN_ALLOW_DISABLED_PARQUET_VECTORIZED_READER.key -> "false") { + val df = spark.read.format("delta").load(path) + + val (_, cometPlan) = checkSparkAnswerAndFallbackReason( + df, + "Native Delta scan is incompatible with " + + "spark.sql.parquet.enableVectorizedReader=false") + + val reasons = new ExtendedExplainInfo().getFallbackReasons(cometPlan) + assert( + !reasons.exists(_.contains("Unsupported file format")), + s"Did not expect core's generic fallback reason among: $reasons") + } + } + } + + test( + "vectorized reader disabled still claims natively when the safety conf allows it " + + "(claim-direction control for the decline above)") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark.range(0, 100).write.format("delta").save(path) + + withSQLConf( + "spark.sql.parquet.enableVectorizedReader" -> "false", + CometConf.COMET_SCAN_ALLOW_DISABLED_PARQUET_VECTORIZED_READER.key -> "true") { + val df = spark.read.format("delta").load(path) + checkSparkAnswer(df) + assert(deltaNativeScans(df).nonEmpty) + } + } + } + + test("_metadata.file_block_start and file_block_length fall back on a split Delta file") { + // When a file is split, Spark keeps a row group in the split holding its midpoint while + // DataFusion keeps it in the split holding its first page, so these per-split values differ. + withSQLConf("spark.sql.files.maxPartitionBytes" -> "4096") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark + .range(0, 5000) + .selectExpr("id", "concat('value_', cast(id as string)) as s") + .coalesce(1) + .write + .format("delta") + .save(path) + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + assert( + spark.read.format("delta").load(path).rdd.getNumPartitions > 1, + "the file has to be split across partitions for this test to mean anything") + } + + val df = spark.read.format("delta").load(path) + for (column <- Seq("file_block_start", "file_block_length")) { + checkSparkAnswerAndFallbackReason( + df.selectExpr("id", "s", s"_metadata.$column"), + s"Native Delta scan does not support metadata column(s) $column") + } + // A filter alone consumes the value too, as does selecting the whole struct. + checkSparkAnswerAndFallbackReason( + df.filter("_metadata.file_block_start > 0").select("id"), + "Native Delta scan does not support metadata column(s) file_block_start") + checkSparkAnswerAndFallbackReason( + df.selectExpr("id", "_metadata"), + "Native Delta scan does not support metadata column(s) " + + "file_block_start, file_block_length") + // The per-file constants do not depend on which split reads a row group. + checkDeltaNativeScanAnswer(df.selectExpr("id", "s", "_metadata.file_path")) + } + } + } + + test("_metadata.file_block_start falls back on a split deletion-vector Delta file") { + // A deletion-vector read puts every _metadata field in the scan output; the plain read stays + // native (see the claim test above) while a query consuming file_block_start falls back. + withSQLConf("spark.sql.files.maxPartitionBytes" -> "4096") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark + .range(0, 5000) + .selectExpr("id", "concat('value_', cast(id as string)) as s") + .coalesce(1) + .write + .format("delta") + .save(path) + spark.sql( + s"ALTER TABLE delta.`$path` SET TBLPROPERTIES ('delta.enableDeletionVectors' = 'true')") + spark.sql(s"DELETE FROM delta.`$path` WHERE id % 7 = 0") + + val df = spark.read.format("delta").load(path) + checkDeltaNativeScanAnswer(df) + val expected = "Native Delta scan does not support metadata column(s) file_block_start" + val (_, cometPlan) = checkSparkAnswerAndFallbackReason( + df.selectExpr("id", "s", "_metadata.file_block_start"), + expected) + // file_block_length is in the scan output too, but nothing above the scan reads it. + val reasons = new ExtendedExplainInfo().getFallbackReasons(cometPlan) + assert( + reasons.exists(r => r.contains(expected) && !r.contains("file_block_length")), + s"expected '$expected' naming only file_block_start among: $reasons") + } + } + } + + test( + "unsupportedSchemes declines an all-viewfs root-path selection (the same helper " + + "declineReason applies to scanExec.relation.location.rootPaths, ahead of the " + + "selected-file gate)") { + val viewfsUri = new URI("viewfs://cluster/table") + // Precondition, mirroring the selected-file scheme tests below: guards against a fail-open + // native build vacuously passing this test. + assert(!CometScanRule.isNativelyReadableScheme(viewfsUri, Set.empty)) + + val schemes = DeltaScanSupport.unsupportedSchemes(Seq(viewfsUri), Set("hdfs")) + assert(schemes == Set("viewfs")) + } + + test( + "unsupportedSchemes flags an opted-in S3-compliant alias scheme that core's gate admits " + + "(the contrib never hands core its alias set)") { + val blobUri = new URI("blob://bucket/table") + // Precondition: core admits the alias once opted in, and only then. + assert(CometScanRule.isNativelyReadableScheme(blobUri, Set("blob"))) + assert(!CometScanRule.isNativelyReadableScheme(blobUri, Set.empty)) + + assert(DeltaScanSupport.unsupportedSchemes(Seq(blobUri), Set("hdfs")) == Set("blob")) + } + + test( + "s3CompliantAliasSchemeReason declines an opted-in alias root path with the explaining " + + "reason, and passes when the scheme is not opted in or is plain s3a") { + val blobUri = new URI("blob://bucket/table") + val s3aUri = new URI("s3a://bucket/table") + + val optedIn = new Configuration(false) + optedIn.set(CometConf.COMET_S3_COMPLIANT_SCHEMES_KEY, " Blob , minio ") + val reason = DeltaScanSupport.s3CompliantAliasSchemeReason(optedIn, Seq(s3aUri, blobUri)) + assert(reason.isDefined) + assert(reason.get.contains("blob")) + assert(reason.get.contains(CometConf.COMET_S3_COMPLIANT_SCHEMES_KEY)) + assert(reason.get.contains("S3AFileSystem")) + assert(DeltaScanSupport.s3CompliantAliasSchemeReason(optedIn, Seq(s3aUri)).isEmpty) + + val notOptedIn = new Configuration(false) + assert(DeltaScanSupport.s3CompliantAliasSchemeReason(notOptedIn, Seq(blobUri)).isEmpty) + } + + test("libhdfsSchemes parses the list exactly like core's scan gate (trim, lowercase, blanks)") { + withSQLConf(CometConf.COMET_LIBHDFS_SCHEMES.key -> " HDFS , viewfs ,, ") { + assert(DeltaScanSupport.libhdfsSchemes == Set("hdfs", "viewfs")) + } + assert(DeltaScanSupport.libhdfsSchemes == Set("hdfs")) + } + + test("unsupportedSchemes passes an all-file: root-path selection (no regression)") { + assert( + DeltaScanSupport + .unsupportedSchemes(Seq(new URI("file:///tmp/table")), Set("hdfs")) + .isEmpty) + } + + test( + "unsupportedSchemes passes a root-path scheme configured as a libhdfs exemption " + + "(exemption honored for the root-path call site too)") { + val viewfsUri = new URI("viewfs://cluster/table") + assert(!CometScanRule.isNativelyReadableScheme(viewfsUri, Set.empty)) + + assert(DeltaScanSupport.unsupportedSchemes(Seq(viewfsUri), Set("viewfs")).isEmpty) + } + + test( + "objectStoreRejectedPathReason declines a root whose path object_store rejects and " + + "stays None for an ordinary path or a libhdfs-exempt scheme") { + val rejected = new URI("file:///tmp/dir%0A/data") + // Precondition, mirroring the scheme tests above: guards against a fail-open native build + // vacuously passing this test. + assert(!CometScanRule.objectStoreAcceptsPath(rejected)) + assert(CometScanRule.objectStoreAcceptsPath(new URI("file:///tmp/table"))) + + val reason = DeltaScanSupport.objectStoreRejectedPathReason(Seq(rejected), Set("hdfs")) + assert(reason.isDefined) + assert(reason.get.contains("cannot open path 'file:///tmp/dir%0A/data'")) + assert(reason.get.contains("object_store rejects it")) + + assert( + DeltaScanSupport + .objectStoreRejectedPathReason(Seq(new URI("file:///tmp/table")), Set("hdfs")) + .isEmpty) + // A libhdfs-routed scheme never reaches object_store's path parser. + assert( + DeltaScanSupport + .objectStoreRejectedPathReason(Seq(new URI("hdfs://nn/dir%0A/data")), Set("hdfs")) + .isEmpty) + // Userinfo is masked in the reason (redactedAuthority's invariant); the path is kept. + val withUserInfo = DeltaScanSupport.objectStoreRejectedPathReason( + Seq(new URI("s3a://key:secret@bucket/dir%0A/data")), + Set("hdfs")) + assert(withUserInfo.exists(_.contains("'s3a://***@bucket/dir%0A/data'")), s"$withUserInfo") + assert(!withUserInfo.exists(_.contains("secret"))) + } + + test( + "objectStoreRejectedPathReason declines a selected file whose basename object_store " + + "rejects even though its parent directory is accepted") { + // CONVERT TO DELTA keeps the source Parquet basenames, so the rejected character can sit in + // the file name itself; a directory-only probe accepts the parent and misses it. Both URIs + // come from Hadoop's Path, as the selected files do, so they render as `file:/...`. + val file = new Path(new Path("file:/tmp/table"), "part-00000\n.snappy.parquet").toUri + val parent = new Path(file).getParent.toUri + assert(file == new URI("file:/tmp/table/part-00000%0A.snappy.parquet"), s"$file") + assert(CometScanRule.objectStoreAcceptsPath(parent), s"directory probe rejected $parent") + assert(!CometScanRule.objectStoreAcceptsPath(file)) + assert(DeltaScanSupport.objectStoreRejectedPathReason(Seq(parent), Set("hdfs")).isEmpty) + + val reason = DeltaScanSupport.objectStoreRejectedPathReason(Seq(file), Set("hdfs")) + assert( + reason.exists( + _.contains("cannot open path 'file:/tmp/table/part-00000%0A.snappy.parquet': " + + "object_store rejects it")), + s"reason: $reason") + } + + test("multiStoreReason declines data files spanning multiple object-store authorities") { + // Same bucket, different keys: one authority, claimable. + assert( + DeltaScanSupport + .multiStoreReason( + Seq(new URI("s3a://bucket/a/part-0.parquet"), new URI("s3a://bucket/b/part-1.parquet"))) + .isEmpty) + + // Distinct buckets: two authorities, must decline (this is the shallow-clone-across- + // buckets-plus-append shape the shared native scan builder cannot route correctly, since + // it resolves the whole scan's ObjectStoreUrl from the first file only). + val reason = DeltaScanSupport.multiStoreReason( + Seq(new URI("s3a://bucket-a/part-0.parquet"), new URI("s3a://bucket-b/part-1.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("multiple object stores")) + assert(reason.get.contains("bucket-a")) + assert(reason.get.contains("bucket-b")) + + // file:// paths never carry an authority (host/port are always empty), so local scans + // across distinct directories are unaffected. + assert( + DeltaScanSupport + .multiStoreReason( + Seq(new URI("file:///tmp/a/part-0.parquet"), new URI("file:///tmp/b/part-1.parquet"))) + .isEmpty) + } + + test( + "multiStoreReason declines cross-container abfss shallow clones (userinfo normalization)") { + // Same storage account, different containers: URI#getHost drops the userinfo entirely, so + // keying the authority on host alone would collapse containerA and containerB into one + // authority and silently claim a cross-container shallow clone. getAuthority (used by + // uriAuthority) keeps the userinfo, so this must decline. + val reason = DeltaScanSupport.multiStoreReason( + Seq( + new URI("abfss://containerA@account.dfs.core.windows.net/a/part-0.parquet"), + new URI("abfss://containerB@account.dfs.core.windows.net/b/part-1.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("multiple object stores")) + + // Same container: one authority, so multiStoreReason itself still passes this shape + // unchanged (this gate was never touched by the userinfo work). But every abfss:// URI here + // carries userinfo (the container) in its authority, so declineReason's earlier-firing + // userInfoBearingAuthorityReason gate now declines this input before multiStoreReason ever + // runs on it -- pinned directly here since multiStoreReason alone can no longer observe the + // difference between this shape and a truly userinfo-free single-authority scan. + val sameContainer = Seq( + new URI("abfss://container@account.dfs.core.windows.net/a/part-0.parquet"), + new URI("abfss://container@account.dfs.core.windows.net/b/part-1.parquet")) + assert(DeltaScanSupport.multiStoreReason(sameContainer).isEmpty) + assert(DeltaScanSupport.userInfoBearingAuthorityReason(sameContainer).isDefined) + } + + test("multiStoreReason declines data files on more than one ABFS store or beside s3a") { + // Core's extractScanObjectStoreOptions needs one ABFS container and account per scan and no + // other scheme beside ABFS. This gate declines all three shapes before convert reaches it. + val shapes = Seq( + "two accounts" -> Seq( + new URI("abfss://data@accountA.dfs.core.windows.net/t/part-0.parquet"), + new URI("abfss://data@accountB.dfs.core.windows.net/t/part-1.parquet")), + "two containers" -> Seq( + new URI("abfss://dataA@account.dfs.core.windows.net/t/part-0.parquet"), + new URI("abfss://dataB@account.dfs.core.windows.net/t/part-1.parquet")), + "abfss and s3a" -> Seq( + new URI("abfss://data@account.dfs.core.windows.net/t/part-0.parquet"), + new URI("s3a://bucket/t/part-1.parquet"))) + shapes.foreach { case (shape, uris) => + val reason = DeltaScanSupport.multiStoreReason(uris) + assert(reason.exists(_.contains("multiple object stores")), s"$shape: $reason") + } + } + + test( + "multiStoreReason declines distinct underscore-bearing GCS buckets " + + "(URI#getHost null-collapse)") { + // `gs://my_bucket` has an underscore reg-name, which URI#getHost cannot parse -- it returns + // null for the WHOLE authority, not just an empty host. Keying uriAuthority on getHost alone + // would make every underscore-bearing bucket normalize to the same "null host" authority + // regardless of which bucket it actually is, so two distinct underscore buckets would + // wrongly collapse into one authority and never decline -- even though the native side + // parses `gs://my_bucket` and `gs://other_bucket` as genuinely different authorities and + // would hard-error on them. getAuthority (used by uriAuthority) returns the raw authority + // text regardless of RFC 3986 conformance, so this must decline instead. + val reason = DeltaScanSupport.multiStoreReason( + Seq( + new URI("gs://my_bucket/a/part-0.parquet"), + new URI("gs://other_bucket/b/part-1.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("multiple object stores")) + + // Same underscore-bearing bucket: one authority, claimable on both the JVM gate and the + // native check (native side asserted in delta_spark_scan.rs's + // same_underscore_host_bucket_files_pass). + assert( + DeltaScanSupport + .multiStoreReason( + Seq( + new URI("gs://my_bucket/a/part-0.parquet"), + new URI("gs://my_bucket/b/part-1.parquet"))) + .isEmpty) + } + + test( + "userInfoBearingAuthorityReason declines a single userinfo-bearing abfss authority " + + "(the behavior change: one container alone is no longer claimable)") { + val reason = DeltaScanSupport.userInfoBearingAuthorityReason( + Seq(new URI("abfss://container@account.dfs.core.windows.net/a/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("userinfo")) + } + + test( + "userInfoBearingAuthorityReason declines two containers on one storage account " + + "(cross-container deletion-vector authority on a single storage account)") { + val reason = DeltaScanSupport.userInfoBearingAuthorityReason( + Seq( + new URI("abfss://source@account.dfs.core.windows.net/a/part-0.parquet"), + new URI("abfss://clone@account.dfs.core.windows.net/_delta_log/dv/deletion_vector.bin"))) + assert(reason.isDefined) + } + + test( + "userInfoBearingAuthorityReason passes s3a data-file and deletion-vector paths " + + "(no regression for the MinIO live suites)") { + // Same bucket: userinfo-free authority, unaffected. + assert( + DeltaScanSupport + .userInfoBearingAuthorityReason( + Seq( + new URI("s3a://bucket/a/part-0.parquet"), + new URI("s3a://bucket/_delta_log/dv/deletion_vector.bin"))) + .isEmpty) + + // Distinct buckets, still no userinfo on either: this gate only inspects userinfo, so it is + // unaffected by multiStoreReason's separate authority-count decline (ported from the deleted + // storeIdentityCollisionReason suite's "passes distinct s3a buckets" case). + assert( + DeltaScanSupport + .userInfoBearingAuthorityReason( + Seq( + new URI("s3a://bucket-a/part-0.parquet"), + new URI("s3a://bucket-b/deletion_vector.bin"))) + .isEmpty) + } + + test("userInfoBearingAuthorityReason passes file:// paths (no authority at all)") { + assert( + DeltaScanSupport + .userInfoBearingAuthorityReason( + Seq(new URI("file:///tmp/a/part-0.parquet"), new URI("file:///tmp/b/part-1.parquet"))) + .isEmpty) + } + + test( + "userInfoBearingAuthorityReason: underscore-bearing GCS bucket passes without userinfo, " + + "declines with it (raw-authority parsing, not URI#getHost)") { + // `gs://my_bucket` has an underscore reg-name that URI#getHost cannot parse (returns null + // for the whole authority); no userinfo either way, so this must pass. + assert( + DeltaScanSupport + .userInfoBearingAuthorityReason(Seq(new URI("gs://my_bucket/a/part-0.parquet"))) + .isEmpty) + + // Same underscore-bearing bucket, now with userinfo: uriUserInfo's raw last-`@` split still + // finds it even though URI#getHost/getUserInfo would return null for this authority. + val reason = DeltaScanSupport.userInfoBearingAuthorityReason( + Seq(new URI("gs://u1@my_bucket/a/part-0.parquet"))) + assert(reason.isDefined) + } + + test("userInfoBearingAuthorityReason passes an hdfs authority with no userinfo") { + assert( + DeltaScanSupport + .userInfoBearingAuthorityReason(Seq(new URI("hdfs://nn:8020/table/part-0.parquet"))) + .isEmpty) + } + + test( + "userInfoBearingAuthorityReason redacts userinfo out of the decline reason (never leaks " + + "embedded credentials)") { + val reason = DeltaScanSupport.userInfoBearingAuthorityReason( + Seq(new URI("s3a://AKIAEXAMPLE:secr3t@bucket/a/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("bucket")) + assert(!reason.get.contains("secr3t")) + assert(!reason.get.contains("AKIAEXAMPLE")) + } + + test( + "unsupportedSelectedSchemeReason declines an all-viewfs selection, naming the scheme and " + + "the selected-file/DV wording") { + val viewfsUri = new URI("viewfs://cluster/table/part-0.parquet") + // Precondition: guards against a fail-open native build vacuously passing this test -- + // isNativelyReadableScheme falls back to TRUE when the native library can't be consulted + // (see its doc), which would make viewfs look natively readable and this test pass for the + // wrong reason regardless of whether the new gate is even wired up correctly. + assert(!CometScanRule.isNativelyReadableScheme(viewfsUri, Set.empty)) + + val reason = DeltaScanSupport.unsupportedSelectedSchemeReason( + Seq(viewfsUri, new URI("viewfs://cluster/table/part-1.parquet")), + Set("hdfs")) + assert(reason.isDefined) + assert(reason.get.contains("viewfs")) + assert(reason.get.contains("data file or deletion vector")) + } + + test( + "unsupportedSelectedSchemeReason declines a mixed file:+viewfs selection with the scheme " + + "reason (pins its ordering ahead of the authority gates)") { + // A supported-scheme file alongside an unsupported-scheme one: this shape ALSO spans + // multiple object-store authorities (multiStoreReason below would decline it too), but + // declineReason places the scheme gate first, so callers must see the scheme reason here, + // not whatever the authority gates would have said about this same input. + val fileUri = new URI("file:///tmp/table/part-0.parquet") + val viewfsUri = new URI("viewfs://cluster/table/part-1.parquet") + assert(!CometScanRule.isNativelyReadableScheme(viewfsUri, Set.empty)) + + val reason = + DeltaScanSupport.unsupportedSelectedSchemeReason(Seq(fileUri, viewfsUri), Set("hdfs")) + assert(reason.isDefined) + assert(reason.get.contains("viewfs")) + // Confirms this input really would ALSO trip multiStoreReason, so the assertion above is + // meaningfully pinning which reason wins under declineReason's ordering, not merely proving + // the scheme gate fires in isolation. + assert(DeltaScanSupport.multiStoreReason(Seq(fileUri, viewfsUri)).isDefined) + } + + test( + "unsupportedSelectedSchemeReason declines a viewfs deletion-vector absolute path even when " + + "every data file is file:// (proves dvUris is part of the gated URI set)") { + val dvUri = new URI("viewfs://cluster/table/_delta_log/dv/deletion_vector.bin") + assert(!CometScanRule.isNativelyReadableScheme(dvUri, Set.empty)) + + val dataFileUris = Seq(new URI("file:///tmp/table/part-0.parquet")) + val reason = + DeltaScanSupport.unsupportedSelectedSchemeReason(dataFileUris :+ dvUri, Set("hdfs")) + assert(reason.isDefined) + assert(reason.get.contains("viewfs")) + } + + test( + "unsupportedSelectedSchemeReason passes all-file: and all-s3a: selections (no regression " + + "for the MinIO live suites)") { + assert( + DeltaScanSupport + .unsupportedSelectedSchemeReason( + Seq( + new URI("file:///tmp/a/part-0.parquet"), + new URI("file:///tmp/b/deletion_vector.bin")), + Set("hdfs")) + .isEmpty) + assert( + DeltaScanSupport + .unsupportedSelectedSchemeReason( + Seq( + new URI("s3a://bucket/a/part-0.parquet"), + new URI("s3a://bucket/_delta_log/dv/deletion_vector.bin")), + Set("hdfs")) + .isEmpty) + } + + test( + "unsupportedSelectedSchemeReason passes an all-viewfs selection when viewfs is configured " + + "as a libhdfs scheme (exemption honored on the new call site)") { + val viewfsUri = new URI("viewfs://cluster/table/part-0.parquet") + assert(!CometScanRule.isNativelyReadableScheme(viewfsUri, Set.empty)) + + assert( + DeltaScanSupport.unsupportedSelectedSchemeReason(Seq(viewfsUri), Set("viewfs")).isEmpty) + } + + test("mergedObjectStoreOptions unions options across every authority without leaking schemes") { + // The merge must reach a DV sidecar living on a different provider than the data files + // (e.g. S3 data + ABFS deletion vector), and must never hand an unrelated provider's + // credentials to a scan that never referenced it. + val hadoopConf = new org.apache.hadoop.conf.Configuration(false) + hadoopConf.set("fs.s3a.access.key", "s3-access-key") + hadoopConf.set("fs.s3a.secret.key", "s3-secret-key") + hadoopConf.set("fs.azure.account.key.acct.dfs.core.windows.net", "azure-account-key") + + val s3Uri = new URI("s3a://bucket/data.parquet") + val abfssUri = new URI("abfss://container@acct.dfs.core.windows.net/dv.bin") + + val merged = + CometDeltaNativeScan.mergedObjectStoreOptions(hadoopConf, Seq(s3Uri, abfssUri)) + assert(merged.get("fs.s3a.access.key").contains("s3-access-key")) + assert(merged.get("fs.s3a.secret.key").contains("s3-secret-key")) + // ABFS options are what Hadoop resolves for the account, not the fs.azure.* namespace. + val abfsOptions = NativeConfig.extractObjectStoreOptions(hadoopConf, abfssUri) + assert(abfsOptions.nonEmpty) + abfsOptions.foreach { case (key, value) => assert(merged.get(key).contains(value), key) } + + // s3-only input must not leak the azure options into the merged map. + val s3Only = CometDeltaNativeScan.mergedObjectStoreOptions(hadoopConf, Seq(s3Uri)) + assert(s3Only.get("fs.s3a.access.key").contains("s3-access-key")) + assert( + !s3Only.keys.exists(key => key.startsWith("fs.azure.") || key.startsWith("comet.azure."))) + } + + test( + "storeUris dedups by authority: one representative URI per (scheme, authority), even " + + "when DV files live at distinct paths on the same authority") { + // No Spark session involved, and deliberately NOT a file:// scan: a local-path test can't + // exercise a DV sidecar on a foreign authority (extractObjectStoreOptions returns an empty + // map for file://), which is exactly the shape that requires unioning object-store options + // across every authority. Hand-build descriptors via Delta's own factory methods instead of + // going through a real scan/claim. + val tableRootPath = new Path("s3a://bucket-root/table") + val firstFileUri = Some(new URI("s3a://bucket-root/table/part-0.parquet")) + + // Path-based ('p') DV on a different authority than the data files / table root. + val foreignDv = DeletionVectorDescriptor + .onDiskWithAbsolutePath("abfss://acct.dfs.core.windows.net/dv1.bin", 40, 4) + // A SECOND, distinct path on the SAME foreign authority as `foreignDv` -- the shape that + // motivates per-authority dedup: before dedup, N deletion-vector files on one external store + // yielded ~N distinct URIs here (each independently walked by mergedObjectStoreOptions); now + // they collapse to a single representative. + val sameAuthoritySecondDv = DeletionVectorDescriptor + .onDiskWithAbsolutePath("abfss://acct.dfs.core.windows.net/dv2.bin", 40, 4) + // UUID-relative ('u') DV: resolves under the table root's authority (s3a/bucket-root), which + // `firstFileUri` already represents -- must not add a second entry for that authority. + val relativeDv = DeletionVectorDescriptor.onDiskWithRelativePath(UUID.randomUUID(), "", 40, 4) + // Inline ('i') DV: no external URI at all; must not be resolved (would throw -- inline + // descriptors fail `absolutePath`'s `isOnDisk` precondition) and must contribute nothing. + val inlineDv = DeletionVectorDescriptor.inlineInLog(Array[Byte](1, 2, 3), 1) + + val uris = CometDeltaNativeScan.storeUris( + Seq(foreignDv, sameAuthoritySecondDv, relativeDv, inlineDv), + tableRootPath, + firstFileUri) + + // Exactly one representative per authority: s3a/bucket-root (firstFileUri wins -- it is + // first in candidate order, ahead of the table root and the relative DV's resolution) and + // abfss/acct.dfs.core.windows.net (foreignDv wins over sameAuthoritySecondDv, the first DV + // seen on that authority). + assert( + uris == Seq(firstFileUri.get, new URI("abfss://acct.dfs.core.windows.net/dv1.bin")), + s"expected exactly one representative URI per authority, got: $uris") + } + + test("storeUris always includes firstFileUri and the table root even with no DV descriptors") { + val tableRootPath = new Path("file:///tmp/table") + val firstFileUri = Some(new URI("file:///tmp/table/part-0.parquet")) + + // firstFileUri and tableRootPath share the same (empty) file:// authority, so the table root + // is deduped away in favor of firstFileUri, which is first in candidate order. + assert( + CometDeltaNativeScan.storeUris(Seq.empty, tableRootPath, firstFileUri) == + Seq(firstFileUri.get)) + + // No first file (e.g. an empty selected-partitions edge case): table root alone, no crash. + assert( + CometDeltaNativeScan.storeUris(Seq.empty, tableRootPath, None) == + Seq(tableRootPath.toUri)) + } + + test("a scan that selects no files keeps the table root's object-store options") { + withTempPath { dir => + val path = dir.getAbsolutePath + spark + .range(0, 10) + .selectExpr("id", "id % 2 AS p") + .write + .format("delta") + .partitionBy("p") + .save(path) + // The scheme list rides along in every extracted map, so a local root's options are not + // empty. + withSQLConf(CometConf.COMET_S3_COMPLIANT_SCHEMES_KEY -> "minio") { + val df = spark.read.format("delta").load(path).where("p = 5") + assert(df.collect().isEmpty) + val scans = deltaNativeScans(df).collect { case s: CometDeltaNativeScanExec => s } + assert(scans.size == 1, s"expected one native Delta scan:\n${df.queryExecution}") + val options = + DeltaSparkScanEnvelope.unpack(scans.head.nativeOp).getCommon.getObjectStoreOptionsMap + assert(options.asScala.toMap == Map(CometConf.COMET_S3_COMPLIANT_SCHEMES_KEY -> "minio")) + } + } + } + + test("user guide documents every native Delta scan config verbatim, with its default") { + // The config table on the user-guide page is hand-maintained: the doc build cannot see + // DeltaSparkConfigProvider with the current module layout, so this is the only check tying + // each entry's key, doc string, and default to the row GenerateDocs would render. + val file = DeltaScanContribSuite + .findRepoFile("docs/source/user-guide/latest/delta.md") + .getOrElse( + fail("Could not locate docs/source/user-guide/latest/delta.md from this checkout; " + + "set -Dcomet.repo.root or run from the repo or module root")) + val source = scala.io.Source.fromFile(file, "UTF-8") + val tableRows = + try source.getLines().filter(_.startsWith("| `")).toList + finally source.close() + // Renders the row GenerateDocs emits for an entry with a plain default and no env var, which + // is every entry today; an entry using either needs the extra text added here as well. + val expectedRows = DeltaScanConf.all.map { conf => + s"| `${conf.key}` | ${conf.doc.trim} | ${conf.defaultValueString} |" + } + expectedRows.foreach { row => + assert( + tableRows.contains(row), + s"Expected ${file.getAbsolutePath} to contain this table row verbatim:\n$row\n" + + s"Rows present:\n${tableRows.mkString("\n")}") + } + val staleRows = tableRows.filterNot(expectedRows.contains) + assert( + staleRows.isEmpty, + s"${file.getAbsolutePath} has table rows matching no entry in DeltaScanConf.all:\n" + + staleRows.mkString("\n")) + } + + test("CometDeltaS3Suite.s3Required arms the hard failure only for 1 or true") { + // Trimmed and case-insensitive so a padded or upper-cased workflow value still counts; + // anything else, including yes and 0, keeps the fail-soft default so a typo cannot arm it. + assert(!CometDeltaS3Suite.s3Required(None)) + Seq("1", "true", "TRUE ", " True").foreach { v => + assert(CometDeltaS3Suite.s3Required(Some(v)), s"'$v' should arm the hard failure") + } + Seq("", " ", "0", "false", "yes", "on", "required", "11").foreach { v => + assert(!CometDeltaS3Suite.s3Required(Some(v)), s"'$v' must not arm the hard failure") + } + } + + test("CometDeltaS3Suite.requiredFailureMessage names the env var and the cause") { + val message = CometDeltaS3Suite.requiredFailureMessage("no Docker daemon is reachable") + assert(message.contains(CometDeltaS3Suite.S3_REQUIRED_ENV)) + assert(message.contains("no Docker daemon is reachable")) + } + + /** + * Builds a real JCEKS keystore backing `hadoop.security.credential.provider.path`, seeded with + * `entries`, and hands `test` a fresh [[Configuration]] already pointed at it (path only -- + * `entries` are NOT mirrored into the plain conf; callers add plain values themselves when a + * case needs them). Uses `CredentialProviderFactory` directly (the real API `Configuration# + * getPassword` reads through), not a hand-rolled keystore, so these tests exercise the actual + * Hadoop credential-provider resolution path rather than a stand-in for it. The store password + * defaults to `"none"` when neither `HADOOP_CREDSTORE_PASSWORD` nor a password file is set in + * the test environment, which is the JCEKS provider's own documented default -- nothing extra + * to configure here. + */ + private def withJceks(entries: Map[String, String])(test: Configuration => Unit): Unit = { + val storeFile = File.createTempFile("comet-delta-creds", ".jceks") + // JavaKeyStoreProvider creates the backing file itself on first flush; a pre-existing empty + // file (createTempFile always creates one) makes it treat the store as an existing, empty + // keystore instead -- harmless either way for JCEKS, but deleting it first keeps this fixture + // honest about what it is actually exercising (provider-created, not merely provider-opened). + storeFile.delete() + val providerPath = "jceks://file" + storeFile.getAbsolutePath + try { + val buildConf = new Configuration(false) + buildConf.set(CredentialProviderFactory.CREDENTIAL_PROVIDER_PATH, providerPath) + val provider = CredentialProviderFactory.getProviders(buildConf).get(0) + entries.foreach { case (alias, value) => + provider.createCredentialEntry(alias, value.toCharArray) + } + provider.flush() + + val testConf = new Configuration(false) + testConf.set(CredentialProviderFactory.CREDENTIAL_PROVIDER_PATH, providerPath) + test(testConf) + } finally { + storeFile.delete() + } + } + + test( + "gcsHadoopOnlyAuthReason declines a gs data file relying only on a Hadoop service-account " + + "keyfile, naming the key but never the value, and matches the scheme case-insensitively") { + val conf = new Configuration(false) + conf.set("fs.gs.auth.service.account.json.keyfile", "/secret/path/svc-key.json") + val reason = + DeltaScanSupport.gcsHadoopOnlyAuthReason(conf, Seq(new URI("GS://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.gs.auth.service.account.json.keyfile")) + assert(reason.get.contains("mybucket")) + assert(!reason.get.contains("/secret/path/svc-key.json")) + } + + test( + "gcsHadoopOnlyAuthReason passes a gs URI when no fs.gs.auth.* key is set " + + "(Application Default Credentials work in both engines)") { + val conf = new Configuration(false) + assert( + DeltaScanSupport + .gcsHadoopOnlyAuthReason(conf, Seq(new URI("gs://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "gcsHadoopOnlyAuthReason does not fire for s3a/file URIs even when fs.gs.auth.* is set " + + "(scheme-scoped)") { + val conf = new Configuration(false) + conf.set("fs.gs.auth.service.account.json.keyfile", "/secret/path/svc-key.json") + assert( + DeltaScanSupport + .gcsHadoopOnlyAuthReason( + conf, + Seq( + new URI("s3a://mybucket/part-0.parquet"), + new URI("file:///tmp/table/part-0.parquet"))) + .isEmpty) + } + + test( + "gcsHadoopOnlyAuthReason declines when local data files are mixed with an absolute gs " + + "deletion-vector sidecar backed only by a Hadoop keyfile") { + val conf = new Configuration(false) + conf.set("fs.gs.auth.service.account.json.keyfile", "/secret/path/svc-key.json") + val reason = DeltaScanSupport.gcsHadoopOnlyAuthReason( + conf, + Seq( + new URI("file:///tmp/table/part-0.parquet"), + new URI("gs://mybucket/_delta_log/deletion_vector_abc123.bin"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.gs.auth.service.account.json.keyfile")) + assert(reason.get.contains("mybucket")) + } + + test( + "gcsHadoopOnlyAuthReason's decline reason names every offending fs.gs.auth.* key but never " + + "any of their configured values") { + val conf = new Configuration(false) + conf.set("fs.gs.auth.service.account.json.keyfile", "/secret/path/svc-key.json") + conf.set("fs.gs.auth.client.id", "super-secret-client-id-xyz") + val reason = + DeltaScanSupport.gcsHadoopOnlyAuthReason(conf, Seq(new URI("gs://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.gs.auth.service.account.json.keyfile")) + assert(reason.get.contains("fs.gs.auth.client.id")) + assert(!reason.get.contains("/secret/path/svc-key.json")) + assert(!reason.get.contains("super-secret-client-id-xyz")) + } + + test( + "gcsHadoopOnlyAuthReason declines a gs data file relying only on the legacy " + + "google.cloud.auth.* connector prefix, naming the key but never the value") { + val conf = new Configuration(false) + conf.set("google.cloud.auth.service.account.json.keyfile", "/secret/path/svc-key.json") + val reason = + DeltaScanSupport.gcsHadoopOnlyAuthReason(conf, Seq(new URI("gs://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("google.cloud.auth.service.account.json.keyfile")) + assert(reason.get.contains("mybucket")) + assert(!reason.get.contains("/secret/path/svc-key.json")) + } + + test( + "gcsHadoopOnlyAuthReason does not fire for s3a/file URIs even when google.cloud.auth.* is " + + "set (scheme-scoped)") { + val conf = new Configuration(false) + conf.set("google.cloud.auth.service.account.json.keyfile", "/secret/path/svc-key.json") + assert( + DeltaScanSupport + .gcsHadoopOnlyAuthReason( + conf, + Seq( + new URI("s3a://mybucket/part-0.parquet"), + new URI("file:///tmp/table/part-0.parquet"))) + .isEmpty) + } + + test( + "gcsHadoopOnlyAuthReason's decline reason names offending keys under both fs.gs.auth. and " + + "google.cloud.auth. but never any of their configured values") { + val conf = new Configuration(false) + conf.set("fs.gs.auth.client.id", "super-secret-client-id-xyz") + conf.set("google.cloud.auth.service.account.json.keyfile", "/secret/path/svc-key.json") + val reason = + DeltaScanSupport.gcsHadoopOnlyAuthReason(conf, Seq(new URI("gs://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.gs.auth.client.id")) + assert(reason.get.contains("google.cloud.auth.service.account.json.keyfile")) + assert(!reason.get.contains("super-secret-client-id-xyz")) + assert(!reason.get.contains("/secret/path/svc-key.json")) + } + + test( + "gcsHadoopOnlyAuthReason declines a gs data file relying only on the deprecated " + + "fs.gs.service.account.auth.keyfile key (reversed word order vs the modern " + + "fs.gs.auth.service.account.* prefix), naming the key but never the value") { + val conf = new Configuration(false) + conf.set("fs.gs.service.account.auth.keyfile", "/secret/path/svc-key.p12") + val reason = + DeltaScanSupport.gcsHadoopOnlyAuthReason(conf, Seq(new URI("gs://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.gs.service.account.auth.keyfile")) + assert(reason.get.contains("mybucket")) + assert(!reason.get.contains("/secret/path/svc-key.p12")) + } + + test( + "gcsHadoopOnlyAuthReason declines a gs data file relying only on the deprecated " + + "fs.gs.service.account.auth.email key, naming the key but never the value") { + val conf = new Configuration(false) + conf.set("fs.gs.service.account.auth.email", "svc@example-project.iam.gserviceaccount.com") + val reason = + DeltaScanSupport.gcsHadoopOnlyAuthReason(conf, Seq(new URI("gs://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.gs.service.account.auth.email")) + assert(!reason.get.contains("svc@example-project.iam.gserviceaccount.com")) + } + + test( + "gcsHadoopOnlyAuthReason declines a gs data file relying only on the deprecated " + + "google.cloud.service.account.auth.keyfile key, naming the key but never the value") { + val conf = new Configuration(false) + conf.set("google.cloud.service.account.auth.keyfile", "/secret/path/svc-key.p12") + val reason = + DeltaScanSupport.gcsHadoopOnlyAuthReason(conf, Seq(new URI("gs://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("google.cloud.service.account.auth.keyfile")) + assert(!reason.get.contains("/secret/path/svc-key.p12")) + } + + test( + "gcsHadoopOnlyAuthReason declines a gs data file relying only on the deprecated " + + "google.cloud.service.account.auth.email key, naming the key but never the value") { + val conf = new Configuration(false) + conf.set( + "google.cloud.service.account.auth.email", + "svc@example-project.iam.gserviceaccount.com") + val reason = + DeltaScanSupport.gcsHadoopOnlyAuthReason(conf, Seq(new URI("gs://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("google.cloud.service.account.auth.email")) + assert(!reason.get.contains("svc@example-project.iam.gserviceaccount.com")) + } + + test( + "gcsHadoopOnlyAuthReason does not fire for s3a/file URIs even when the deprecated " + + "fs.gs.service.account.auth.* prefix is set (scheme-scoped)") { + val conf = new Configuration(false) + conf.set("fs.gs.service.account.auth.keyfile", "/secret/path/svc-key.p12") + assert( + DeltaScanSupport + .gcsHadoopOnlyAuthReason( + conf, + Seq( + new URI("s3a://mybucket/part-0.parquet"), + new URI("file:///tmp/table/part-0.parquet"))) + .isEmpty) + } + + test( + "gcsHadoopOnlyAuthReason declines on fs.gs.auth.type, a suffix no fixed prefix list ever " + + "enumerated (predicate-based matching instead of a prefix table)") { + val conf = new Configuration(false) + conf.set("fs.gs.auth.type", "SERVICE_ACCOUNT_JSON_KEYFILE") + val reason = + DeltaScanSupport.gcsHadoopOnlyAuthReason(conf, Seq(new URI("gs://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.gs.auth.type")) + assert(reason.get.contains("mybucket")) + assert(!reason.get.contains("SERVICE_ACCOUNT_JSON_KEYFILE")) + } + + test( + "gcsHadoopOnlyAuthReason declines on fs.gs.auth.client.id, naming the key but never the " + + "value") { + val conf = new Configuration(false) + conf.set("fs.gs.auth.client.id", "super-secret-client-id-xyz") + val reason = + DeltaScanSupport.gcsHadoopOnlyAuthReason(conf, Seq(new URI("gs://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.gs.auth.client.id")) + assert(!reason.get.contains("super-secret-client-id-xyz")) + } + + test( + "s3ConfigDivergenceReason declines when access/secret keys exist only in a JCEKS " + + "keystore, naming the base key and bucket but never the secret") { + withJceks(Map("fs.s3a.access.key" -> "AKIAEXAMPLE", "fs.s3a.secret.key" -> "s3cr3tValue")) { + conf => + val reason = DeltaScanSupport.s3ConfigDivergenceReason( + conf, + Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.access.key")) + assert(reason.get.contains("mybucket")) + assert(!reason.get.contains("AKIAEXAMPLE")) + assert(!reason.get.contains("s3cr3tValue")) + } + } + + test( + "s3ConfigDivergenceReason passes when only plain keys are set and no provider path is " + + "configured (zero-I/O precheck exit)") { + val conf = new Configuration(false) + conf.set("fs.s3a.access.key", "AKIAPLAIN") + conf.set("fs.s3a.secret.key", "plainSecret") + assert( + DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "s3ConfigDivergenceReason passes when the provider path is set and the plain keys match " + + "the keystore (plain keys consistent with the credential provider)") { + withJceks(Map("fs.s3a.access.key" -> "AKIAMATCH", "fs.s3a.secret.key" -> "matchingSecret")) { + conf => + conf.set("fs.s3a.access.key", "AKIAMATCH") + conf.set("fs.s3a.secret.key", "matchingSecret") + assert( + DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + } + + test( + "s3ConfigDivergenceReason declines when the keystore value differs from a shadowed plain " + + "value") { + withJceks(Map("fs.s3a.access.key" -> "AKIAKEYSTORE")) { conf => + conf.set("fs.s3a.access.key", "AKIADIFFERENTPLAIN") + val reason = DeltaScanSupport.s3ConfigDivergenceReason( + conf, + Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.access.key")) + assert(!reason.get.contains("AKIAKEYSTORE")) + } + } + + test( + "s3ConfigDivergenceReason declines on an S3A-scoped provider path immediately, without " + + "touching a nonexistent keystore (Arm A proves no keystore I/O)") { + val tempDir = Files.createTempDirectory("comet-delta-no-keystore") + try { + val conf = new Configuration(false) + val nonexistentPath = "jceks://file" + tempDir + "/does-not-exist.jceks" + conf.set("fs.s3a.security.credential.provider.path", nonexistentPath) + // No exception from a missing file is the point of this test: Arm A declines on the + // presence of the S3A-scoped path key alone, never reading it. + val reason = DeltaScanSupport.s3ConfigDivergenceReason( + conf, + Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.security.credential.provider.path")) + } finally { + Files.delete(tempDir) + } + } + + test( + "s3ConfigDivergenceReason passes file:// URIs regardless of any provider path " + + "(S3-only scope)") { + val conf = new Configuration(false) + conf.set("hadoop.security.credential.provider.path", "jceks://file/nonexistent.jceks") + assert( + DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("file:///tmp/table/part-0.parquet"))) + .isEmpty) + } + + test( + "s3ConfigDivergenceReason declines via a per-bucket credential alias " + + "(fs.s3a.bucket.mybucket.access.key)") { + withJceks(Map("fs.s3a.bucket.mybucket.access.key" -> "AKIABUCKETSCOPED")) { conf => + val reason = DeltaScanSupport.s3ConfigDivergenceReason( + conf, + Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.access.key")) + assert(reason.get.contains("mybucket")) + assert(!reason.get.contains("AKIABUCKETSCOPED")) + } + } + + test( + "s3ConfigDivergenceReason declines via a long-form per-bucket credential alias " + + "(fs.s3a.bucket.mybucket.fs.s3a.access.key), a Hadoop S3AUtils.lookupPassword alias " + + "the short-form check alone misses") { + withJceks( + Map( + "fs.s3a.bucket.mybucket.fs.s3a.access.key" -> "AKIALONGFORM", + "fs.s3a.bucket.mybucket.fs.s3a.secret.key" -> "longFormSecret")) { conf => + val reason = DeltaScanSupport.s3ConfigDivergenceReason( + conf, + Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.access.key")) + assert(reason.get.contains("mybucket")) + assert(!reason.get.contains("AKIALONGFORM")) + assert(!reason.get.contains("longFormSecret")) + } + } + + test( + "s3ConfigDivergenceReason declines via a long-form per-bucket credential alias even when " + + "different plain global keys are also configured (Hadoop would resolve the long-form " + + "keystore value first; native reads only the differing plain globals)") { + withJceks( + Map( + "fs.s3a.bucket.mybucket.fs.s3a.access.key" -> "AKIALONGFORM", + "fs.s3a.bucket.mybucket.fs.s3a.secret.key" -> "longFormSecret")) { conf => + conf.set("fs.s3a.access.key", "AKIADIFFERENTGLOBAL") + conf.set("fs.s3a.secret.key", "differentGlobalSecret") + val reason = DeltaScanSupport.s3ConfigDivergenceReason( + conf, + Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.access.key")) + assert(reason.get.contains("mybucket")) + assert(!reason.get.contains("AKIALONGFORM")) + assert(!reason.get.contains("longFormSecret")) + assert(!reason.get.contains("AKIADIFFERENTGLOBAL")) + assert(!reason.get.contains("differentGlobalSecret")) + } + } + + test( + "s3ConfigDivergenceReason declines on a long-form per-bucket provider path immediately, " + + "without touching a nonexistent keystore (Arm A proves no keystore I/O)") { + val tempDir = Files.createTempDirectory("comet-delta-no-keystore-long-bucket") + try { + val conf = new Configuration(false) + val nonexistentPath = "jceks://file" + tempDir + "/does-not-exist.jceks" + conf.set("fs.s3a.bucket.mybucket.fs.s3a.security.credential.provider.path", nonexistentPath) + // No exception from a missing file is the point of this test: Arm A declines on the + // presence of the long-form bucket-scoped path key alone, never reading it. + val reason = DeltaScanSupport.s3ConfigDivergenceReason( + conf, + Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert( + reason.get.contains("fs.s3a.bucket.mybucket.fs.s3a.security.credential.provider.path")) + } finally { + Files.delete(tempDir) + } + } + + test( + "s3ConfigDivergenceReason passes when only plain global keys are set and no provider " + + "path is configured, including the long-form bucket provider path (control: unaffected " + + "by the new long-form aliases)") { + val conf = new Configuration(false) + conf.set("fs.s3a.access.key", "AKIAPLAIN") + conf.set("fs.s3a.secret.key", "plainSecret") + assert( + DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "s3ConfigDivergenceReason declines without throwing when the keystore is " + + "corrupt/unreadable (global arm try/catch containment)") { + val corruptFile = File.createTempFile("comet-delta-corrupt-creds", ".jceks") + try { + Files.write(corruptFile.toPath, Array[Byte](1, 2, 3, 4, 5, 6, 7, 8)) + val conf = new Configuration(false) + conf.set( + "hadoop.security.credential.provider.path", + "jceks://file" + corruptFile.getAbsolutePath) + // Must not throw: a corrupt/unreadable keystore must decline this bucket, not escape and + // abort planning for the whole session. + val reason = DeltaScanSupport.s3ConfigDivergenceReason( + conf, + Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + } finally { + corruptFile.delete() + } + } + + test( + "s3ConfigDivergenceReason declines when a plain long-form bucket credential key is set " + + "with nothing else (Hadoop resolves it, native's short-then-global lookup never sees " + + "it), naming the base key and bucket but never a credential value") { + val conf = new Configuration(false) + conf.set("fs.s3a.bucket.mybucket.fs.s3a.access.key", "AKIALONGPLAIN") + conf.set("fs.s3a.bucket.mybucket.fs.s3a.secret.key", "longPlainSecret") + val reason = DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.access.key")) + assert(reason.get.contains("mybucket")) + assert(!reason.get.contains("AKIALONGPLAIN")) + assert(!reason.get.contains("longPlainSecret")) + } + + test("s3ConfigDivergenceReason declines when a long-form bucket credential holds a Hadoop " + + s"$${...} reference that DOES resolve, with nothing else set (substitution alone does not " + + "erase the long-form divergence: native's short-then-global read never consults the long " + + "form regardless of what it expands to), naming the base key and bucket but never a value") { + val conf = new Configuration(false) + conf.set("review.longFormAccess", "AKIALONGRESOLVED") + conf.set("fs.s3a.bucket.mybucket.fs.s3a.access.key", s"$${review.longFormAccess}") + val reason = DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.access.key")) + assert(reason.get.contains("mybucket")) + assert(!reason.get.contains("AKIALONGRESOLVED")) + assert(!reason.get.contains(s"$${review.longFormAccess}")) + } + + test( + "s3ConfigDivergenceReason declines when a plain long-form bucket credential diverges " + + "from a different plain global value (Hadoop would use the long-form bucket value; " + + "native would use the differing global)") { + val conf = new Configuration(false) + conf.set("fs.s3a.bucket.mybucket.fs.s3a.access.key", "AKIALONGPLAIN") + conf.set("fs.s3a.access.key", "AKIADIFFERENTGLOBALPLAIN") + val reason = DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.access.key")) + assert(reason.get.contains("mybucket")) + assert(!reason.get.contains("AKIALONGPLAIN")) + assert(!reason.get.contains("AKIADIFFERENTGLOBALPLAIN")) + } + + test( + "s3ConfigDivergenceReason declines when the plain long-form and short-form bucket " + + "credential keys are set to DIFFERENT values (Hadoop's SimpleAWSCredentialsProvider " + + "resolves the long pair; native resolves the short pair, so they diverge), naming the " + + "base key and bucket but never a credential value") { + val conf = new Configuration(false) + conf.set("fs.s3a.bucket.mybucket.fs.s3a.access.key", "long-ak") + conf.set("fs.s3a.bucket.mybucket.fs.s3a.secret.key", "long-sk") + conf.set("fs.s3a.bucket.mybucket.access.key", "short-ak") + conf.set("fs.s3a.bucket.mybucket.secret.key", "short-sk") + val reason = DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.access.key")) + assert(reason.get.contains("mybucket")) + assert(!reason.get.contains("long-ak")) + assert(!reason.get.contains("short-ak")) + } + + test( + "control: S3AUtils#propagateBucketOptions folds a long-form bucket option into the " + + "unread key fs.s3a.fs.s3a.endpoint, proving Hadoop itself ignores the long form for " + + "general (non-credential) per-bucket options -- unlike lookupPassword for credentials, " + + "no Comet gate exists (or is needed) for this case") { + val conf = new Configuration(false) + conf.set("fs.s3a.bucket.mybucket.fs.s3a.endpoint", "long-form.example.com") + conf.set("fs.s3a.endpoint", "global.example.com") + + // Real Hadoop code, not a Comet stand-in: S3AFileSystem#initialize assigns exactly this + // result to the `conf` it reads ENDPOINT/PATH_STYLE_ACCESS/etc. from. + val propagated = S3AUtils.propagateBucketOptions(conf, "mybucket") + assert(propagated.get("fs.s3a.endpoint") == "global.example.com") + assert(propagated.get("fs.s3a.fs.s3a.endpoint") == "long-form.example.com") + } + + test( + "s3ConfigDivergenceReason declines when a bucket-scoped credential references another " + + "bucket-scoped key that Hadoop's real propagate-then-resolve order shadows the global " + + "value with (Hadoop resolves the bucket-scoped referent; native, which never propagates " + + "bucket options, still resolves the global one), naming the base key and bucket but " + + "never a credential value") { + val conf = new Configuration(false) + conf.set("fs.s3a.bucket.mybucket.access.key", s"$${fs.s3a.custom.ref}") + conf.set("fs.s3a.bucket.mybucket.custom.ref", "bucket-scoped-value") + conf.set("fs.s3a.custom.ref", "global-value") + + // Real Hadoop code, not a Comet stand-in: this is exactly what S3AFileSystem#initialize + // assigns to the `conf` it later reads fs.s3a.access.key from -- the bucket-scoped + // fs.s3a.bucket.mybucket.custom.ref overwrites the global fs.s3a.custom.ref BEFORE the + // ${...} reference in the propagated fs.s3a.bucket.mybucket.access.key is ever substituted. + val propagated = S3AUtils.propagateBucketOptions(conf, "mybucket") + assert(propagated.get("fs.s3a.custom.ref") == "bucket-scoped-value") + assert(propagated.get("fs.s3a.bucket.mybucket.access.key") == "bucket-scoped-value") + + val reason = DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.access.key")) + assert(reason.get.contains("mybucket")) + assert(!reason.get.contains("bucket-scoped-value")) + assert(!reason.get.contains("global-value")) + } + + test( + "s3ConfigDivergenceReason passes when a bucket-scoped credential references another " + + "bucket-scoped key whose propagated value happens to equal the global value (no actual " + + "divergence, despite the same shadowing mechanism as the declining case above)") { + val conf = new Configuration(false) + conf.set("fs.s3a.bucket.mybucket.access.key", s"$${fs.s3a.custom.ref}") + conf.set("fs.s3a.bucket.mybucket.custom.ref", "same-value") + conf.set("fs.s3a.custom.ref", "same-value") + assert( + DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "s3ConfigDivergenceReason declines, without any keystore I/O, when a bucket-scoped " + + s"long-form credential-provider-path key (Arm A) is itself set via a $${...} reference " + + "to another bucket-scoped key that Hadoop's real propagate-then-resolve order shadows " + + "the global value with -- naming only the provider path key and bucket, never either " + + "resolved path") { + // Uses the LONG form (fs.s3a.bucket.B.fs.s3a.security.credential.provider.path), not the + // short form, deliberately: propagateBucketOptions folds ANY fs.s3a.bucket.B. key into + // a global fs.s3a. key. For the short form, is + // "security.credential.provider.path", so it propagates into the GLOBAL S3A-scoped provider + // path key itself (fs.s3a.security.credential.provider.path) -- correctly triggering the + // OTHER Arm A branch instead, since real Hadoop would see the same thing. The long form's + // is "fs.s3a.security.credential.provider.path", which propagates into the inert, + // double-prefixed fs.s3a.fs.s3a.security.credential.provider.path key instead, isolating + // the long-form bucket-scoped branch this test targets. + val conf = new Configuration(false) + conf.set( + "fs.s3a.bucket.mybucket.fs.s3a.security.credential.provider.path", + s"$${fs.s3a.custom.ref}") + conf.set( + "fs.s3a.bucket.mybucket.custom.ref", + "jceks://file/does-not-exist-bucket-scoped.jceks") + conf.set("fs.s3a.custom.ref", "jceks://file/does-not-exist-global.jceks") + + // Real Hadoop code, not a Comet stand-in: this is exactly what S3AFileSystem#initialize + // assigns to the `conf` it later reads the bucket-scoped provider path from -- the + // bucket-scoped fs.s3a.bucket.mybucket.custom.ref overwrites the global fs.s3a.custom.ref + // BEFORE the ${...} reference in the propagated provider path key is ever substituted. + val propagated = S3AUtils.propagateBucketOptions(conf, "mybucket") + assert( + propagated.get("fs.s3a.custom.ref") == "jceks://file/does-not-exist-bucket-scoped.jceks") + assert( + propagated.get("fs.s3a.bucket.mybucket.fs.s3a.security.credential.provider.path") == + "jceks://file/does-not-exist-bucket-scoped.jceks") + // Confirms the long form's propagated target is the inert double-prefixed key, NOT the + // global S3A-scoped provider path key -- i.e. this test genuinely isolates the long-form + // bucket-scoped branch rather than accidentally exercising the global-S3A-path branch. + assert(propagated.get("fs.s3a.security.credential.provider.path") == null) + + // Neither referenced path exists on disk -- if this gate mistakenly tried to open either + // as a keystore instead of declining on the key's mere presence (Arm A), it would throw + // rather than return a reason, which this test would catch. + val reason = DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.bucket.mybucket.fs.s3a.security.credential.provider.path")) + assert(reason.get.contains("mybucket")) + assert(!reason.get.contains("does-not-exist-bucket-scoped")) + assert(!reason.get.contains("does-not-exist-global")) + } + + test( + "s3ConfigDivergenceReason declines when the long-form and global bucket credential keys " + + "share the same value but the short-form bucket keys are set to EMPTY strings (Hadoop's " + + "SimpleAWSCredentialsProvider resolves the long pair via lookupPassword's skip-empty " + + "semantics; native's get_config_trimmed resolves the short pair's mere PRESENCE, landing " + + "on empty credentials instead), naming the base key and bucket but never a credential " + + "value") { + val conf = new Configuration(false) + conf.set("fs.s3a.bucket.mybucket.fs.s3a.access.key", "shared-ak") + conf.set("fs.s3a.bucket.mybucket.fs.s3a.secret.key", "shared-sk") + conf.set("fs.s3a.access.key", "shared-ak") + conf.set("fs.s3a.secret.key", "shared-sk") + conf.set("fs.s3a.bucket.mybucket.access.key", "") + conf.set("fs.s3a.bucket.mybucket.secret.key", "") + val reason = DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.access.key")) + assert(reason.get.contains("mybucket")) + assert(!reason.get.contains("shared-ak")) + } + + test("s3ConfigDivergenceReason declines when the short-form bucket credential keys hold only " + + "whitespace: native's get_config_trimmed still resolves the key's mere PRESENCE before " + + "trimming its value, so a whitespace-only short-form key diverges from Hadoop's long-form " + + "resolution exactly like an outright empty one") { + val conf = new Configuration(false) + conf.set("fs.s3a.bucket.mybucket.fs.s3a.access.key", "shared-ak") + conf.set("fs.s3a.bucket.mybucket.fs.s3a.secret.key", "shared-sk") + conf.set("fs.s3a.access.key", "shared-ak") + conf.set("fs.s3a.secret.key", "shared-sk") + conf.set("fs.s3a.bucket.mybucket.access.key", " ") + conf.set("fs.s3a.bucket.mybucket.secret.key", " ") + val reason = DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.access.key")) + assert(reason.get.contains("mybucket")) + } + + test( + "control: s3ConfigDivergenceReason passes when the short-form bucket credential keys are " + + "absent rather than empty, so Hadoop's long-form resolution and native's short-then-global " + + "resolution both land on the same shared pair") { + val conf = new Configuration(false) + conf.set("fs.s3a.bucket.mybucket.fs.s3a.access.key", "shared-ak") + conf.set("fs.s3a.bucket.mybucket.fs.s3a.secret.key", "shared-sk") + conf.set("fs.s3a.access.key", "shared-ak") + conf.set("fs.s3a.secret.key", "shared-sk") + assert( + DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "control: s3ConfigDivergenceReason passes when a global-only value (no bucket override at " + + "all, so Hadoop's and native's effective values are the exact same conf entry) carries " + + "incidental leading/trailing whitespace, such as Hadoop's own multi-line " + + "fs.s3a.aws.credentials.provider default -- trimming must apply symmetrically to both " + + "sides of the comparison, or an untouched default value would diverge from itself and " + + "decline every S3 scan") { + val conf = new Configuration(false) + conf.set( + "fs.s3a.aws.credentials.provider", + "\n org.apache.hadoop.fs.s3a.TemporaryAWSCredentialsProvider,\n " + + "org.apache.hadoop.fs.s3a.SimpleAWSCredentialsProvider\n ") + assert( + DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + s"s3ConfigDivergenceReason declines when a per-bucket override redirects the $${...} " + + "reference inside the short-form bucket endpoint while a long-form endpoint alias holds " + + "the value native resolves (Hadoop's endpoint consumer is propagateBucketOptions plus " + + "plain Configuration#get, which follows the redirected reference and never reads the " + + "long form at all)") { + val conf = new Configuration(false) + conf.set("fs.s3a.custom.ref", "https://store-a.example") + conf.set("fs.s3a.bucket.data-bucket.custom.ref", "https://store-b.example") + conf.set("fs.s3a.bucket.data-bucket.endpoint", s"$${fs.s3a.custom.ref}") + conf.set("fs.s3a.bucket.data-bucket.fs.s3a.endpoint", "https://store-a.example") + + // Real Hadoop code, not a Comet stand-in: propagation overwrites the global referent with + // the per-bucket custom.ref BEFORE the endpoint's ${...} reference is substituted, so + // Hadoop's plain endpoint read lands on store-b -- while native, which never propagates, + // expands the same reference against the original conf and lands on store-a. + val propagated = S3AUtils.propagateBucketOptions(conf, "data-bucket") + assert(propagated.get("fs.s3a.endpoint") == "https://store-b.example") + + val reason = DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://data-bucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.endpoint")) + assert(reason.get.contains("data-bucket")) + assert(!reason.get.contains("store-a")) + assert(!reason.get.contains("store-b")) + } + + test( + "control: s3ConfigDivergenceReason passes the same endpoint shape without the per-bucket " + + s"referent override (the $${...} reference expands identically with and without " + + "bucket-option propagation, so Hadoop's plain-get endpoint read and native agree)") { + val conf = new Configuration(false) + conf.set("fs.s3a.custom.ref", "https://store-a.example") + conf.set("fs.s3a.bucket.data-bucket.endpoint", s"$${fs.s3a.custom.ref}") + conf.set("fs.s3a.bucket.data-bucket.fs.s3a.endpoint", "https://store-a.example") + assert( + DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://data-bucket/part-0.parquet"))) + .isEmpty) + } + + test( + "s3ConfigDivergenceReason passes when ONLY the long-form bucket endpoint alias is set: " + + "propagateBucketOptions folds it into the unread fs.s3a.fs.s3a.endpoint key, so Hadoop's " + + "plain endpoint read and native's short-then-global read both resolve nothing") { + val conf = new Configuration(false) + conf.set("fs.s3a.bucket.data-bucket.fs.s3a.endpoint", "https://store-a.example") + assert( + DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://data-bucket/part-0.parquet"))) + .isEmpty) + } + + test( + "s3ConfigDivergenceReason declines an obsolete plaintext credential pair when " + + "hadoop.security.credential.clear-text-fallback is false: Configuration#getPassword " + + "ignores plain conf then, so Hadoop's SimpleAWSCredentialsProvider reports no " + + "credentials and the chain proceeds to the environment -- while native would sign every " + + "request with the stale static pair") { + val conf = new Configuration(false) + conf.set("fs.s3a.access.key", "stale-ak") + conf.set("fs.s3a.secret.key", "stale-sk") + conf.set( + "fs.s3a.aws.credentials.provider", + "org.apache.hadoop.fs.s3a.SimpleAWSCredentialsProvider," + + "com.amazonaws.auth.EnvironmentVariableCredentialsProvider") + conf.set("hadoop.security.credential.clear-text-fallback", "false") + val reason = DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.access.key")) + assert(reason.get.contains("mybucket")) + assert(!reason.get.contains("stale-ak")) + assert(!reason.get.contains("stale-sk")) + } + + test( + "control: s3ConfigDivergenceReason passes the same plaintext pair and provider chain when " + + "clear-text-fallback keeps its default (true): getPassword falls back to plain conf, so " + + "Hadoop and native resolve the identical static pair") { + val conf = new Configuration(false) + conf.set("fs.s3a.access.key", "stale-ak") + conf.set("fs.s3a.secret.key", "stale-sk") + conf.set( + "fs.s3a.aws.credentials.provider", + "org.apache.hadoop.fs.s3a.SimpleAWSCredentialsProvider," + + "com.amazonaws.auth.EnvironmentVariableCredentialsProvider") + assert( + DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "s3ConfigDivergenceReason passes with clear-text-fallback=false when no plaintext " + + "credential is set anywhere: both sides resolve no credentials, and the provider-class " + + "key itself stays comparable (its consumer is Configuration#getClasses, plain conf, " + + "which the fallback flag never gates)") { + val conf = new Configuration(false) + conf.set( + "fs.s3a.aws.credentials.provider", + "org.apache.hadoop.fs.s3a.SimpleAWSCredentialsProvider," + + "com.amazonaws.auth.EnvironmentVariableCredentialsProvider") + conf.set("hadoop.security.credential.clear-text-fallback", "false") + assert( + DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + // ----------------------------------------------------------------------------------------- + // providerClassGateReason: declines an unsupported credential-provider class (or + // invalid combination) before the scan is claimed, rather than letting it fail during + // execution in s3.rs's build_aws_credential_provider_metadata. + // ----------------------------------------------------------------------------------------- + + private val nativeSupportedProviderClasses = Seq( + "org.apache.hadoop.fs.s3a.auth.IAMInstanceCredentialsProvider", + "org.apache.hadoop.fs.s3a.SimpleAWSCredentialsProvider", + "org.apache.hadoop.fs.s3a.TemporaryAWSCredentialsProvider", + "org.apache.hadoop.fs.s3a.AnonymousAWSCredentialsProvider", + "software.amazon.awssdk.auth.credentials.ContainerCredentialsProvider", + "com.amazonaws.auth.ContainerCredentialsProvider", + "com.amazonaws.auth.EC2ContainerCredentialsProviderWrapper", + "software.amazon.awssdk.auth.credentials.InstanceProfileCredentialsProvider", + "com.amazonaws.auth.InstanceProfileCredentialsProvider", + "software.amazon.awssdk.auth.credentials.EnvironmentVariableCredentialsProvider", + "com.amazonaws.auth.EnvironmentVariableCredentialsProvider", + "software.amazon.awssdk.auth.credentials.WebIdentityTokenFileCredentialsProvider", + "com.amazonaws.auth.WebIdentityTokenCredentialsProvider", + "software.amazon.awssdk.auth.credentials.ProfileCredentialsProvider", + "com.amazonaws.auth.profile.ProfileCredentialsProvider", + "software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider", + "com.amazonaws.auth.AnonymousAWSCredentials") + + test( + "providerClassGateReason passes when aws.credentials.provider is unset (native's " + + "default AWS SDK provider chain)") { + val conf = new Configuration(false) + assert( + DeltaScanSupport + .providerClassGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test("providerClassGateReason passes for every credential provider class s3.rs supports") { + nativeSupportedProviderClasses.foreach { className => + val conf = new Configuration(false) + conf.set("fs.s3a.aws.credentials.provider", className) + val reason = + DeltaScanSupport + .providerClassGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isEmpty, s"Expected $className to be claimable, but got: $reason") + } + } + + test( + "providerClassGateReason declines an unsupported credential provider class, naming the " + + "key and the bucket but never the configured value") { + val conf = new Configuration(false) + conf.set("fs.s3a.aws.credentials.provider", "com.example.CustomCredentialsProvider") + val reason = + DeltaScanSupport + .providerClassGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.aws.credentials.provider")) + assert(!reason.get.contains("com.example.CustomCredentialsProvider")) + assert(reason.get.contains("mybucket")) + } + + test( + "providerClassGateReason declines via the per-bucket short form, honoring bucket-scoped " + + "override (mirrors get_config's short-then-global resolution)") { + val conf = new Configuration(false) + conf.set( + "fs.s3a.aws.credentials.provider", + "org.apache.hadoop.fs.s3a.SimpleAWSCredentialsProvider") + conf.set( + "fs.s3a.bucket.mybucket.aws.credentials.provider", + "com.example.CustomCredentialsProvider") + val reason = + DeltaScanSupport + .providerClassGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(!reason.get.contains("com.example.CustomCredentialsProvider")) + } + + test( + "providerClassGateReason passes a comma-separated list of entirely supported provider " + + "classes (native chains them via build_chained_aws_credential_provider_metadata)") { + val conf = new Configuration(false) + conf.set( + "fs.s3a.aws.credentials.provider", + "org.apache.hadoop.fs.s3a.SimpleAWSCredentialsProvider, " + + "software.amazon.awssdk.auth.credentials.EnvironmentVariableCredentialsProvider") + assert( + DeltaScanSupport + .providerClassGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "providerClassGateReason declines a comma-separated list containing one unsupported " + + "class, naming neither class") { + val conf = new Configuration(false) + conf.set( + "fs.s3a.aws.credentials.provider", + "org.apache.hadoop.fs.s3a.SimpleAWSCredentialsProvider,com.example.Bogus") + val reason = + DeltaScanSupport + .providerClassGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(!reason.get.contains("com.example.Bogus")) + assert(!reason.get.contains("SimpleAWSCredentialsProvider")) + } + + test( + "providerClassGateReason declines an anonymous provider mixed with another provider " + + "(native's build_credential_provider rejects this combination at execution time)") { + val conf = new Configuration(false) + conf.set( + "fs.s3a.aws.credentials.provider", + "org.apache.hadoop.fs.s3a.AnonymousAWSCredentialsProvider," + + "org.apache.hadoop.fs.s3a.SimpleAWSCredentialsProvider") + val reason = + DeltaScanSupport + .providerClassGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("anonymous")) + } + + test( + "providerClassGateReason passes a solo anonymous provider (native returns None -- an " + + "unsigned client -- rather than erroring; only a MIX with other providers is rejected)") { + val conf = new Configuration(false) + conf.set( + "fs.s3a.aws.credentials.provider", + "org.apache.hadoop.fs.s3a.AnonymousAWSCredentialsProvider") + assert( + DeltaScanSupport + .providerClassGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "providerClassGateReason passes AssumedRoleCredentialProvider with an unset " + + "assumed.role.credentials.provider (native defaults to its own always-supported " + + "[Simple, EnvironmentVariable] fallback)") { + val conf = new Configuration(false) + conf.set( + "fs.s3a.aws.credentials.provider", + "org.apache.hadoop.fs.s3a.auth.AssumedRoleCredentialProvider") + assert( + DeltaScanSupport + .providerClassGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "providerClassGateReason declines AssumedRoleCredentialProvider whose " + + "assumed.role.credentials.provider names an unsupported base provider class") { + val conf = new Configuration(false) + conf.set( + "fs.s3a.aws.credentials.provider", + "org.apache.hadoop.fs.s3a.auth.AssumedRoleCredentialProvider") + conf.set("fs.s3a.assumed.role.credentials.provider", "com.example.BogusBaseProvider") + val reason = + DeltaScanSupport + .providerClassGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(!reason.get.contains("com.example.BogusBaseProvider")) + assert(reason.get.contains("fs.s3a.assumed.role.credentials.provider")) + } + + test( + "providerClassGateReason declines AssumedRoleCredentialProvider whose " + + "assumed.role.credentials.provider names an anonymous base provider (native rejects ANY " + + "anonymous entry here, not just a mix)") { + val conf = new Configuration(false) + conf.set( + "fs.s3a.aws.credentials.provider", + "org.apache.hadoop.fs.s3a.auth.AssumedRoleCredentialProvider") + conf.set( + "fs.s3a.assumed.role.credentials.provider", + "org.apache.hadoop.fs.s3a.AnonymousAWSCredentialsProvider") + val reason = + DeltaScanSupport + .providerClassGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("anonymous")) + } + + test( + "assumedRolePolicyGateReason declines when a global assumed-role session policy is " + + "configured") { + val conf = new Configuration(false) + conf.set("fs.s3a.assumed.role.policy", """{"Version":"2012-10-17","Statement":[]}""") + val reason = + DeltaScanSupport + .assumedRolePolicyGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.assumed.role.policy")) + // The policy document itself is security-sensitive configuration; never leak it. + assert(!reason.get.contains("2012-10-17")) + } + + test( + "assumedRolePolicyGateReason declines when a bucket-scoped assumed-role session policy " + + "is configured") { + val conf = new Configuration(false) + conf.set( + "fs.s3a.bucket.mybucket.assumed.role.policy", + """{"Version":"2012-10-17","Statement":[]}""") + val reason = + DeltaScanSupport + .assumedRolePolicyGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.assumed.role.policy")) + assert(!reason.get.contains("2012-10-17")) + } + + test("assumedRolePolicyGateReason admits when no assumed-role session policy is configured") { + val conf = new Configuration(false) + conf.set("fs.s3a.assumed.role.arn", "arn:aws:iam::123456789012:role/reader") + assert( + DeltaScanSupport + .assumedRolePolicyGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "providerClassGateReason ignores assumed.role.credentials.provider when " + + "AssumedRoleCredentialProvider is not itself in play (dead config on the native side)") { + val conf = new Configuration(false) + conf.set( + "fs.s3a.aws.credentials.provider", + "org.apache.hadoop.fs.s3a.SimpleAWSCredentialsProvider") + conf.set("fs.s3a.assumed.role.credentials.provider", "com.example.BogusBaseProvider") + assert( + DeltaScanSupport + .providerClassGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "providerClassGateReason passes when the global aws.credentials.provider key holds a " + + "Hadoop variable reference that Configuration#get expands to a supported class " + + "(post-substitution, native's plain-conf extraction sees the same expanded class name " + + "the class-support check does, so no divergence exists to decline)") { + val conf = new Configuration(false) + conf.set("review.provider", "org.apache.hadoop.fs.s3a.SimpleAWSCredentialsProvider") + conf.set("fs.s3a.aws.credentials.provider", s"$${review.provider}") + assert( + DeltaScanSupport + .providerClassGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "providerClassGateReason passes when a bucket-scoped short-form " + + "aws.credentials.provider override holds a variable reference that expands to a " + + "supported class, even though the global key is a different supported literal") { + val conf = new Configuration(false) + conf.set( + "fs.s3a.aws.credentials.provider", + "org.apache.hadoop.fs.s3a.SimpleAWSCredentialsProvider") + conf.set("review.bucketProvider", "org.apache.hadoop.fs.s3a.SimpleAWSCredentialsProvider") + conf.set("fs.s3a.bucket.mybucket.aws.credentials.provider", s"$${review.bucketProvider}") + assert( + DeltaScanSupport + .providerClassGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "providerClassGateReason passes when the assumed-role base-provider key holds a " + + "variable reference that expands to a supported base class") { + val conf = new Configuration(false) + conf.set( + "fs.s3a.aws.credentials.provider", + "org.apache.hadoop.fs.s3a.auth.AssumedRoleCredentialProvider") + conf.set("review.baseProvider", "org.apache.hadoop.fs.s3a.SimpleAWSCredentialsProvider") + conf.set("fs.s3a.assumed.role.credentials.provider", s"$${review.baseProvider}") + assert( + DeltaScanSupport + .providerClassGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "providerClassGateReason passes literal provider classes with no variable references " + + "(unaffected by variable expansion, still runs the class-support gate)") { + val conf = new Configuration(false) + conf.set( + "fs.s3a.aws.credentials.provider", + "org.apache.hadoop.fs.s3a.SimpleAWSCredentialsProvider") + assert( + DeltaScanSupport + .providerClassGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + + val badConf = new Configuration(false) + badConf.set("fs.s3a.aws.credentials.provider", "com.example.CustomCredentialsProvider") + val reason = + DeltaScanSupport + .providerClassGateReason(badConf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(!reason.get.contains("com.example.CustomCredentialsProvider")) + } + + test( + "providerClassGateReason declines rather than throws when a two-key mutual Hadoop " + + "variable-reference cycle involves fs.s3a.aws.credentials.provider, called DIRECTLY " + + "(not routed through s3ConfigDivergenceReason, which masks this for the same keys when " + + "checked first -- this pins the gate's OWN containment, not that coupling)") { + val conf = new Configuration(false) + conf.set("fs.s3a.aws.credentials.provider", s"$${fs.s3a.assumed.role.credentials.provider}") + conf.set("fs.s3a.assumed.role.credentials.provider", s"$${fs.s3a.aws.credentials.provider}") + val reason = + DeltaScanSupport + .providerClassGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.aws.credentials.provider")) + assert(reason.get.contains("IllegalStateException")) + assert(!reason.get.contains(s"$${fs.s3a.assumed.role.credentials.provider}")) + assert(!reason.get.contains(s"$${fs.s3a.aws.credentials.provider}")) + } + + test( + "s3ConfigDivergenceReason passes when both the plain long-form and short-form bucket " + + "credential keys are set to the EQUAL value (Hadoop's long-first resolution and " + + "native's short-then-global resolution agree)") { + val conf = new Configuration(false) + conf.set("fs.s3a.bucket.mybucket.fs.s3a.access.key", "AKIASAMEBOTH") + conf.set("fs.s3a.bucket.mybucket.access.key", "AKIASAMEBOTH") + assert( + DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "s3ConfigDivergenceReason passes when the plain long-form bucket credential value equals " + + "the plain global value (both sides resolve to the same value)") { + val conf = new Configuration(false) + conf.set("fs.s3a.bucket.mybucket.fs.s3a.access.key", "AKIASAME") + conf.set("fs.s3a.access.key", "AKIASAME") + assert( + DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "s3ConfigDivergenceReason passes with only a plain short-form bucket credential key set " + + "(control: unaffected by the long-form plain-value check)") { + val conf = new Configuration(false) + conf.set("fs.s3a.bucket.mybucket.access.key", "AKIASHORTONLY") + assert( + DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "s3ConfigDivergenceReason passes when a credential key holds a Hadoop variable reference " + + "that Configuration#get expands to a literal (post-substitution, native's plain-conf " + + "extraction forwards the SAME expanded value this comparator reads, so both sides agree)") { + val conf = new Configuration(false) + conf.set("review.access", "AKIAEXAMPLE") + conf.set("fs.s3a.access.key", s"$${review.access}") + assert( + DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "s3ConfigDivergenceReason passes when credential keys hold literal values with no " + + "variable references") { + val conf = new Configuration(false) + conf.set("fs.s3a.access.key", "AKIALITERAL") + conf.set("fs.s3a.secret.key", "literalSecretValue") + assert( + DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "s3ConfigDivergenceReason passes when a credential key references an undefined variable " + + "(Hadoop leaves the literal unresolved, so native and Hadoop see the identical value)") { + val conf = new Configuration(false) + conf.set("fs.s3a.access.key", s"$${undefined.var}") + assert( + DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "s3ConfigDivergenceReason passes when a bucket-scoped short-form credential alias holds " + + "a variable reference that expands identically for both sides (alias-set coverage " + + "beyond the plain global key)") { + val conf = new Configuration(false) + conf.set("review.secret", "topSecretValue") + conf.set("fs.s3a.bucket.mybucket.secret.key", s"$${review.secret}") + assert( + DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "s3ConfigDivergenceReason does not throw for a credential key that is its own Hadoop " + + "variable reference (Configuration#get's substitution loop converges immediately -- " + + "the raw and expanded literals are already equal -- so this is the same safe shape as " + + "an undefined variable, not a MAX_SUBST failure)") { + val conf = new Configuration(false) + conf.set("fs.s3a.secret.key", "realSecretValue") + conf.set("fs.s3a.access.key", s"$${fs.s3a.access.key}") + val reason = DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert( + reason.isEmpty, + "expected no decline (and no exception) for a literal " + + s"self-reference, since it resolves to the same unexpanded text on both sides: $reason") + } + + test( + "s3ConfigDivergenceReason declines rather than throws when two credential keys form a " + + "mutual Hadoop variable-reference cycle (Configuration#get raises IllegalStateException " + + s"once $${...} substitution recurses past Hadoop's MAX_SUBST bound)") { + val conf = new Configuration(false) + conf.set("fs.s3a.access.key", s"$${fs.s3a.secret.key}") + conf.set("fs.s3a.secret.key", s"$${fs.s3a.access.key}") + val reason = DeltaScanSupport + .s3ConfigDivergenceReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("IllegalStateException")) + assert(!reason.get.contains("realSecretValue")) + } + + test( + "unsupportedEncryptionAlgorithmReason declines a bucket configured with global SSE-C, " + + "naming the algorithm key and the algorithm but never the customer-provided key value") { + val conf = new Configuration(false) + conf.set("fs.s3a.encryption.algorithm", "SSE-C") + conf.set("fs.s3a.encryption.key", "c3VwZXItc2VjcmV0LWN1c3RvbWVyLWtleQ==") + val reason = DeltaScanSupport + .unsupportedEncryptionAlgorithmReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.encryption.algorithm")) + assert(reason.get.contains("SSE-C")) + assert(reason.get.contains("mybucket")) + assert(!reason.get.contains("c3VwZXItc2VjcmV0LWN1c3RvbWVyLWtleQ==")) + } + + test( + "unsupportedEncryptionAlgorithmReason matches the SSE-C algorithm value " + + "case-insensitively, mirroring S3AEncryptionMethods#getMethod's equalsIgnoreCase " + + "parsing") { + val conf = new Configuration(false) + conf.set("fs.s3a.encryption.algorithm", "sse-c") + val reason = DeltaScanSupport + .unsupportedEncryptionAlgorithmReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("mybucket")) + } + + test( + "unsupportedEncryptionAlgorithmReason declines a bucket configured with the deprecated " + + "fs.s3a.server-side-encryption-algorithm spelling of SSE-C, naming the algorithm key " + + "actually consulted but never the customer-provided key value") { + val conf = new Configuration(false) + conf.set("fs.s3a.server-side-encryption-algorithm", "SSE-C") + conf.set("fs.s3a.server-side-encryption.key", "c3VwZXItc2VjcmV0LWN1c3RvbWVyLWtleQ==") + val reason = DeltaScanSupport + .unsupportedEncryptionAlgorithmReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + // Names EITHER spelling, never both/neither: hadoop-aws's S3AFileSystem statically registers + // this exact pair as a Configuration-level deprecated alias (verified via javap -- + // S3AFileSystem.addDeprecatedKeys() calls Configuration.addDeprecations, a field static on + // Hadoop's Configuration class, process-wide once S3AFileSystem's class has loaded anywhere + // in this JVM -- which a real Spark job has always done by the time it evaluates this gate, + // since reading the S3 table at all requires that class). Once registered, + // Configuration#get resolves either literal key to the same value transparently, so which + // name THIS gate happens to read the value under depends on whether that static + // registration already ran elsewhere in the test JVM, not on anything this test controls. + assert( + reason.get.contains("fs.s3a.server-side-encryption-algorithm") || + reason.get.contains("fs.s3a.encryption.algorithm")) + assert(reason.get.contains("mybucket")) + assert(!reason.get.contains("c3VwZXItc2VjcmV0LWN1c3RvbWVyLWtleQ==")) + } + + test( + "unsupportedEncryptionAlgorithmReason declines only the bucket whose per-bucket " + + "SHORT-form key sets SSE-C, leaving an unrelated bucket unaffected") { + val conf = new Configuration(false) + conf.set("fs.s3a.bucket.secure-bucket.encryption.algorithm", "SSE-C") + val declined = DeltaScanSupport.unsupportedEncryptionAlgorithmReason( + conf, + Seq(new URI("s3a://secure-bucket/part-0.parquet"))) + assert(declined.isDefined) + assert(declined.get.contains("secure-bucket")) + + assert( + DeltaScanSupport + .unsupportedEncryptionAlgorithmReason( + conf, + Seq(new URI("s3a://other-bucket/part-0.parquet"))) + .isEmpty) + } + + test( + "unsupportedEncryptionAlgorithmReason DOES fire for SSE-C set only via the LONG " + + "per-bucket form: S3AUtils#lookupBucketSecret is long-then-short, " + + "decompiled from hadoop-aws 3.3.4's S3AUtils.class -- unlike a plain propagated option, " + + "the encryption algorithm's bucket tier DOES consult fs.s3a.bucket.B.fs.s3a.encryption." + + "algorithm, and Hadoop's own reader picks SSE-C from it, so this must decline exactly " + + "like the short-form case above") { + val conf = new Configuration(false) + conf.set("fs.s3a.bucket.secure-bucket.fs.s3a.encryption.algorithm", "SSE-C") + val reason = DeltaScanSupport.unsupportedEncryptionAlgorithmReason( + conf, + Seq(new URI("s3a://secure-bucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("secure-bucket")) + assert(reason.get.contains("SSE-C")) + } + + private val EncryptionAlgorithmKeys = + Seq("fs.s3a.encryption.algorithm", "fs.s3a.server-side-encryption-algorithm") + + private val DataBucketUris = Seq(new URI("s3a://data-bucket/part-0.parquet")) + + /** The algorithm Hadoop's own S3AFileSystem selects for `bucket`: propagate, then resolve. */ + private def hadoopEncryptionAlgorithm(conf: Configuration, bucket: String): String = { + // S3AFileSystem's static initializer registers the deprecated algorithm key as an alias of + // the canonical one; a real filesystem always has it loaded before it resolves anything. + // The alias is process-wide, so afterwards setting the deprecated key also stores the + // canonical one and the deprecated-key iterations mostly re-test the canonical path. + new S3AFileSystem() + S3AUtils + .getEncryptionAlgorithm(bucket, S3AUtils.propagateBucketOptions(conf, bucket)) + .getMethod + } + + /** An algorithm key set to `${fs.s3a.custom.ref}`, with a global and optional bucket ref. */ + private def referencedAlgorithmConf( + algorithmKey: String, + globalRef: String, + bucketRef: Option[String]): Configuration = { + val conf = new Configuration(false) + conf.set("fs.s3a.custom.ref", globalRef) + bucketRef.foreach(conf.set("fs.s3a.bucket.data-bucket.custom.ref", _)) + conf.set(algorithmKey, s"$${fs.s3a.custom.ref}") + conf + } + + test( + "unsupportedEncryptionAlgorithmReason declines SSE-C that a bucket override selects " + + s"through a $${...} reference, resolving the algorithm on the propagated view as Hadoop " + + "does, and never names the customer key") { + val customerKey = "c3VwZXItc2VjcmV0LWN1c3RvbWVyLWtleQ==" + for (algorithmKey <- EncryptionAlgorithmKeys) { + val conf = referencedAlgorithmConf(algorithmKey, "AES256", Some("SSE-C")) + conf.set("fs.s3a.encryption.key", customerKey) + assert(hadoopEncryptionAlgorithm(conf, "data-bucket") == "SSE-C") + + val reason = DeltaScanSupport.unsupportedEncryptionAlgorithmReason(conf, DataBucketUris) + assert(reason.isDefined, s"expected a decline for $algorithmKey") + assert(reason.get.contains("SSE-C")) + assert(reason.get.contains("data-bucket")) + assert(!reason.get.contains(customerKey)) + } + } + + test( + s"unsupportedEncryptionAlgorithmReason passes AES256 selected through a $${...} reference " + + "when no bucket override redirects it") { + for (algorithmKey <- EncryptionAlgorithmKeys) { + val conf = referencedAlgorithmConf(algorithmKey, "AES256", None) + assert(hadoopEncryptionAlgorithm(conf, "data-bucket") == "AES256") + assert( + DeltaScanSupport.unsupportedEncryptionAlgorithmReason(conf, DataBucketUris).isEmpty, + s"expected $algorithmKey to pass") + } + } + + test( + "unsupportedEncryptionAlgorithmReason passes SSE-KMS that a bucket override selects " + + s"through a $${...} reference (both views resolve an allowlisted algorithm)") { + for (algorithmKey <- EncryptionAlgorithmKeys) { + val conf = referencedAlgorithmConf(algorithmKey, "AES256", Some("SSE-KMS")) + assert(hadoopEncryptionAlgorithm(conf, "data-bucket") == "SSE-KMS") + assert( + DeltaScanSupport.unsupportedEncryptionAlgorithmReason(conf, DataBucketUris).isEmpty, + s"expected $algorithmKey to pass") + } + } + + test( + "unsupportedEncryptionAlgorithmReason declines when the unpropagated view resolves SSE-C " + + "even though the bucket override makes Hadoop select AES256 (a disagreement only " + + "declines)") { + for (algorithmKey <- EncryptionAlgorithmKeys) { + val conf = referencedAlgorithmConf(algorithmKey, "SSE-C", Some("AES256")) + assert(hadoopEncryptionAlgorithm(conf, "data-bucket") == "AES256") + val reason = DeltaScanSupport.unsupportedEncryptionAlgorithmReason(conf, DataBucketUris) + assert(reason.isDefined, s"expected a decline for $algorithmKey") + assert(reason.get.contains("SSE-C")) + } + } + + test( + "unsupportedEncryptionAlgorithmReason never prints an algorithm value that expands to the " + + "customer key, whether the propagated or the raw view resolves it") { + val customerKey = "TESTKEY-NEVER-PRINT-123" + val keyRef = s"$${fs.s3a.encryption.key}" + // (global ref, bucket ref): both views, propagated view only, raw view only. + val shapes = Seq((keyRef, None), ("AES256", Some(keyRef)), (keyRef, Some("AES256"))) + for (algorithmKey <- EncryptionAlgorithmKeys; (globalRef, bucketRef) <- shapes) { + val conf = referencedAlgorithmConf(algorithmKey, globalRef, bucketRef) + conf.set("fs.s3a.encryption.key", customerKey) + val reason = DeltaScanSupport.unsupportedEncryptionAlgorithmReason(conf, DataBucketUris) + assert(reason.isDefined, s"expected a decline for $algorithmKey, $globalRef, $bucketRef") + assert(reason.get.contains("not a recognised S3 encryption algorithm")) + assert(!reason.get.contains(customerKey)) + } + } + + test("providerClassGateReason never prints a provider value that expands to a credential") { + val secret = "TESTKEY-NEVER-PRINT-123" + val providerKeys = + Seq("fs.s3a.aws.credentials.provider", "fs.s3a.bucket.mybucket.aws.credentials.provider") + for (key <- providerKeys) { + val conf = new Configuration(false) + conf.set("fs.s3a.secret.key", secret) + conf.set(key, s"$${fs.s3a.secret.key}") + val reason = DeltaScanSupport + .providerClassGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined, s"expected a decline for $key") + assert(!reason.get.contains(secret)) + } + } + + test( + "unsupportedEncryptionAlgorithmReason declines CSE-KMS (client-side encryption): the " + + "native Parquet reader has no client-side decryption layer, so it would read raw " + + "ciphertext where Hadoop's own reader, which decrypts client-side via the SDK, succeeds") { + val conf = new Configuration(false) + conf.set("fs.s3a.encryption.algorithm", "CSE-KMS") + val reason = DeltaScanSupport + .unsupportedEncryptionAlgorithmReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.encryption.algorithm")) + assert(reason.get.contains("CSE-KMS")) + assert(reason.get.contains("mybucket")) + } + + test( + "unsupportedEncryptionAlgorithmReason declines CSE-CUSTOM (client-side encryption) the " + + "same way as CSE-KMS") { + val conf = new Configuration(false) + conf.set("fs.s3a.encryption.algorithm", "CSE-CUSTOM") + val reason = DeltaScanSupport + .unsupportedEncryptionAlgorithmReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("CSE-CUSTOM")) + } + + test( + "unsupportedEncryptionAlgorithmReason declines an unrecognized future algorithm string " + + "(allowlist semantics: anything not positively confirmed transparent declines, rather " + + "than a blocklist that would silently admit a new Hadoop encryption method)") { + val conf = new Configuration(false) + conf.set("fs.s3a.encryption.algorithm", "SOME-FUTURE-ALGORITHM") + val reason = DeltaScanSupport + .unsupportedEncryptionAlgorithmReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + // Only a value from the fixed list of known algorithm names is ever printed. + assert(reason.get.contains("not a recognised S3 encryption algorithm")) + assert(!reason.get.contains("SOME-FUTURE-ALGORITHM")) + } + + test( + "unsupportedEncryptionAlgorithmReason passes for AES256, SSE-KMS, DSSE-KMS, and for no " + + "encryption configured at all (S3 decrypts these server-side algorithms transparently " + + "on GET/HEAD given read permission alone; only SSE-C requires a client-sent key, and " + + "only CSE-* requires client-side decryption)") { + for (algorithm <- Seq("AES256", "SSE-KMS", "DSSE-KMS")) { + val conf = new Configuration(false) + conf.set("fs.s3a.encryption.algorithm", algorithm) + conf.set("fs.s3a.encryption.key", "arn:aws:kms:us-east-1:123456789012:key/abc-123") + assert( + DeltaScanSupport + .unsupportedEncryptionAlgorithmReason( + conf, + Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty, + s"expected $algorithm to be allowlisted") + } + + val unsetConf = new Configuration(false) + assert( + DeltaScanSupport + .unsupportedEncryptionAlgorithmReason( + unsetConf, + Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "unsupportedEncryptionAlgorithmReason declines when the algorithm is stored ONLY in a " + + "JCEKS keystore as SSE-C, naming the algorithm key and value but never any keystore " + + "material (buildEncryptionSecrets resolves the algorithm via getPassword, which this " + + "gate now mirrors instead of the JCEKS-blind plain-conf read that used to under-decline " + + "this case)") { + withJceks(Map("fs.s3a.encryption.algorithm" -> "SSE-C")) { conf => + val reason = DeltaScanSupport + .unsupportedEncryptionAlgorithmReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.encryption.algorithm")) + assert(reason.get.contains("SSE-C")) + assert(reason.get.contains("mybucket")) + } + } + + test( + "unsupportedEncryptionAlgorithmReason passes a plaintext-only SSE-C algorithm when " + + "clear-text-fallback is false and no credential provider is configured (getPassword " + + "masks the plaintext value, so Hadoop's own buildEncryptionSecrets resolves NO algorithm " + + "and issues plain GETs with no SSE-C key header -- exactly what native issues)") { + // Admitting is safe on the shape's own terms: S3 enforces the customer-key-header + // requirement at the protocol level against every reader, so on a genuinely SSE-C-encrypted + // object both engines fail loudly and identically (400, no header sent), and on an + // unencrypted object both read the same bytes. No config state here lets Hadoop decrypt + // while native reads ciphertext. + val conf = new Configuration(false) + conf.set("fs.s3a.encryption.algorithm", "SSE-C") + conf.set("hadoop.security.credential.clear-text-fallback", "false") + assert( + DeltaScanSupport + .unsupportedEncryptionAlgorithmReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "unsupportedEncryptionAlgorithmReason passes when the algorithm is stored ONLY in a JCEKS " + + "keystore as AES256 (allowlisted even through the keystore-aware resolution path)") { + withJceks(Map("fs.s3a.encryption.algorithm" -> "AES256")) { conf => + assert(DeltaScanSupport + .unsupportedEncryptionAlgorithmReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + } + + test( + "unsupportedEncryptionAlgorithmReason declines without throwing when the keystore backing " + + "the algorithm is corrupt/unreadable (global arm try/catch containment, same pattern as " + + "s3ConfigDivergenceReason's corrupt-keystore test)") { + val corruptFile = File.createTempFile("comet-delta-corrupt-encryption-creds", ".jceks") + try { + Files.write(corruptFile.toPath, Array[Byte](1, 2, 3, 4, 5, 6, 7, 8)) + val conf = new Configuration(false) + conf.set( + "hadoop.security.credential.provider.path", + "jceks://file" + corruptFile.getAbsolutePath) + // Must not throw: a corrupt/unreadable keystore must decline this bucket, not escape and + // abort planning for the whole session. + val reason = DeltaScanSupport + .unsupportedEncryptionAlgorithmReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + } finally { + corruptFile.delete() + } + } + + test( + "unsupportedEncryptionAlgorithmReason declines on an S3A-scoped provider path immediately " + + "when resolving the algorithm, without touching a nonexistent keystore (Arm A proves no " + + "keystore I/O), even though no algorithm key is set in plain conf") { + val tempDir = Files.createTempDirectory("comet-delta-encryption-no-keystore") + try { + val conf = new Configuration(false) + val nonexistentPath = "jceks://file" + tempDir + "/does-not-exist.jceks" + conf.set("fs.s3a.security.credential.provider.path", nonexistentPath) + val reason = DeltaScanSupport + .unsupportedEncryptionAlgorithmReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.security.credential.provider.path")) + } finally { + Files.delete(tempDir) + } + } + + test( + "unsupportedEncryptionAlgorithmReason declines a global provider path with any entry that " + + "is not a local keystore, without opening it") { + val expected = "Native Delta scan cannot verify the S3 settings for data-bucket " + + "(hadoop.security.credential.provider.path names a credential keystore that is not on " + + "the local filesystem, and the settings cannot be verified without reading it)" + val recorder = classOf[OpenRecordingFileSystem].getName + // (provider path, whether fs.file.impl names a non-default class) + val shapes = Seq( + // A scheme of any name can map to S3AFileSystem; only the class decides for Hadoop. + ("jceks://s3x@keys/s3.jceks", false), + ("jceks://s3a@keys-bucket/s3.jceks", false), + ("jceks://file/tmp/local.jceks,jceks://s3@keys-bucket/s3.jceks", false), + ("jceks://hdfs@namenode/keys.jceks", false), + ("jceks://viewfs@cluster/keys.jceks", false), + ("not a uri", false), + // Hadoop resolves the nested file scheme through fs.file.impl too. + ("jceks://file/etc/k.jceks", true), + // A third-party provider factory would run its own code during planning. + ("vault://file/k", false)) + val failures = shapes.flatMap { case (providerPath, overrideFileImpl) => + OpenRecordingFileSystem.opened.set(false) + val conf = new Configuration(false) + conf.set("hadoop.security.credential.provider.path", providerPath) + val schemes = Seq("s3", "s3a", "s3x", "hdfs", "viewfs") ++ + (if (overrideFileImpl) Seq("file") else Nil) + schemes.foreach { scheme => + conf.set(s"fs.$scheme.impl", recorder) + conf.set(s"fs.$scheme.impl.disable.cache", "true") + } + conf.set("fs.s3a.encryption.algorithm", "SSE-C") + val reason = DeltaScanSupport.unsupportedEncryptionAlgorithmReason(conf, DataBucketUris) + if (OpenRecordingFileSystem.opened.get()) { + Some(s"$providerPath was opened") + } else if (!reason.contains(expected)) { + Some(s"unexpected reason for $providerPath: $reason") + } else { + None + } + } + assert(failures.isEmpty, failures.mkString("\n")) + } + + test( + "unsupportedEncryptionAlgorithmReason still evaluates a local keystore or user:/// " + + "provider path, falling back to the plaintext algorithm") { + for (providerPath <- Seq("jceks://file/tmp/does-not-exist.jceks", "user:///")) { + val passing = new Configuration(false) + passing.set("hadoop.security.credential.provider.path", providerPath) + passing.set("fs.s3a.encryption.algorithm", "AES256") + assert( + DeltaScanSupport.unsupportedEncryptionAlgorithmReason(passing, DataBucketUris).isEmpty, + s"expected $providerPath with AES256 to pass") + + val declining = new Configuration(false) + declining.set("hadoop.security.credential.provider.path", providerPath) + declining.set("fs.s3a.encryption.algorithm", "SSE-C") + val reason = + DeltaScanSupport.unsupportedEncryptionAlgorithmReason(declining, DataBucketUris) + assert(reason.exists(_.contains("fs.s3a.encryption.algorithm=SSE-C")), s"got $reason") + } + } + + test( + "unsupportedEncryptionAlgorithmReason does not fire for non-S3 URIs even when SSE-C is " + + "configured globally (scheme-scoped, no S3 bucket to derive from a file:// or gs:// " + + "URI)") { + val conf = new Configuration(false) + conf.set("fs.s3a.encryption.algorithm", "SSE-C") + conf.set("fs.s3a.encryption.key", "c3VwZXItc2VjcmV0LWN1c3RvbWVyLWtleQ==") + assert( + DeltaScanSupport + .unsupportedEncryptionAlgorithmReason( + conf, + Seq( + new URI("file:///tmp/table/part-0.parquet"), + new URI("gs://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "hadoopOnlyEndpointGateReason declines a scheme-less fs.s3a.endpoint with SSL disabled, " + + "since Hadoop addresses it over http:// and native assumes https://") { + val conf = new Configuration(false) + conf.set("fs.s3a.endpoint", "minio:9000") + conf.set("fs.s3a.connection.ssl.enabled", "false") + val reason = DeltaScanSupport.hadoopOnlyEndpointGateReason( + conf, + Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.connection.ssl.enabled")) + assert(reason.get.contains("mybucket")) + } + + test("hadoopOnlyEndpointGateReason admits a scheme-less endpoint with SSL at its default") { + val conf = new Configuration(false) + conf.set("fs.s3a.endpoint", "minio:9000") + assert( + DeltaScanSupport + .hadoopOnlyEndpointGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "hadoopOnlyEndpointGateReason admits an endpoint that carries its own scheme even with " + + "SSL disabled, since Hadoop does not rewrite it") { + val conf = new Configuration(false) + conf.set("fs.s3a.endpoint", "http://minio:9000") + conf.set("fs.s3a.connection.ssl.enabled", "false") + assert( + DeltaScanSupport + .hadoopOnlyEndpointGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "hadoopOnlyEndpointGateReason declines via a short-form per-bucket " + + "fs.s3a.bucket.mybucket.connection.ssl.enabled=false") { + val conf = new Configuration(false) + conf.set("fs.s3a.endpoint", "minio:9000") + conf.set("fs.s3a.bucket.mybucket.connection.ssl.enabled", "false") + val reason = DeltaScanSupport.hadoopOnlyEndpointGateReason( + conf, + Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("mybucket")) + assert( + DeltaScanSupport + .hadoopOnlyEndpointGateReason(conf, Seq(new URI("s3a://otherbucket/part-0.parquet"))) + .isEmpty) + } + + test("hadoopOnlyEndpointGateReason declines when an assumed-role STS endpoint is set") { + Seq("fs.s3a.assumed.role.sts.endpoint", "fs.s3a.assumed.role.sts.endpoint.region").foreach { + key => + val conf = new Configuration(false) + conf.set(key, "sts.eu-west-1.amazonaws.com") + val reason = DeltaScanSupport.hadoopOnlyEndpointGateReason( + conf, + Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined, key) + assert(reason.get.contains(key)) + assert(reason.get.contains("mybucket")) + } + } + + test( + "hadoopOnlyEndpointGateReason declines a short-form per-bucket assumed-role STS endpoint " + + "and admits when nothing endpoint-related is configured") { + val conf = new Configuration(false) + conf.set("fs.s3a.bucket.mybucket.assumed.role.sts.endpoint", "sts.eu-west-1.amazonaws.com") + assert( + DeltaScanSupport + .hadoopOnlyEndpointGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isDefined) + assert( + DeltaScanSupport + .hadoopOnlyEndpointGateReason( + new Configuration(false), + Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "proxyGateReason declines a bucket configured with a global fs.s3a.proxy.host, naming the " + + "key and bucket but never any proxy credential") { + val conf = new Configuration(false) + conf.set("fs.s3a.proxy.host", "proxy.internal.example.com") + conf.set("fs.s3a.proxy.port", "8080") + conf.set("fs.s3a.proxy.username", "proxyuser") + conf.set("fs.s3a.proxy.password", "proxySecretValue") + val reason = + DeltaScanSupport.proxyGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.proxy.host")) + assert(reason.get.contains("mybucket")) + assert(!reason.get.contains("proxyuser")) + assert(!reason.get.contains("proxySecretValue")) + assert(!reason.get.contains("proxy.internal.example.com")) + } + + test( + "proxyGateReason declines via a short-form per-bucket fs.s3a.proxy.host " + + "(fs.s3a.bucket.mybucket.proxy.host)") { + val conf = new Configuration(false) + conf.set("fs.s3a.bucket.mybucket.proxy.host", "proxy.internal.example.com") + val reason = + DeltaScanSupport.proxyGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.proxy.host")) + assert(reason.get.contains("mybucket")) + } + + test( + "proxyGateReason passes on a lone long-form per-bucket fs.s3a.proxy.host " + + "(fs.s3a.bucket.mybucket.fs.s3a.proxy.host): propagateBucketOptions folds it into the " + + "unread key fs.s3a.fs.s3a.proxy.host, and the host's real consumer is a plain getTrimmed " + + "on the propagated conf that never checks any long-form alias") { + // S3AUtils#initProxySupport (hadoop-aws 3.3.4) and AWSClientConfig#createProxyConfiguration + // (3.4.x) both read the host as conf.getTrimmed("fs.s3a.proxy.host", ""), so a lone long + // alias never routes Hadoop through a proxy, same fold as the fs.s3a.endpoint control above. + val conf = new Configuration(false) + conf.set("fs.s3a.bucket.mybucket.fs.s3a.proxy.host", "proxy.internal.example.com") + assert( + DeltaScanSupport + .proxyGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "proxyGateReason passes when no fs.s3a.proxy.host is configured anywhere (zero-I/O, no " + + "provider path set)") { + val conf = new Configuration(false) + assert( + DeltaScanSupport + .proxyGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "proxyGateReason declines only the bucket whose proxy host is actually configured, " + + "leaving an unrelated bucket unaffected") { + val conf = new Configuration(false) + conf.set("fs.s3a.bucket.proxied-bucket.proxy.host", "proxy.internal.example.com") + val declined = + DeltaScanSupport.proxyGateReason(conf, Seq(new URI("s3a://proxied-bucket/part-0.parquet"))) + assert(declined.isDefined) + assert(declined.get.contains("proxied-bucket")) + + assert( + DeltaScanSupport + .proxyGateReason(conf, Seq(new URI("s3a://other-bucket/part-0.parquet"))) + .isEmpty) + } + + test( + "proxyGateReason does not fire for non-S3 URIs even when fs.s3a.proxy.host is configured " + + "globally (scheme-scoped, no S3 bucket to derive from a file:// or gs:// URI)") { + val conf = new Configuration(false) + conf.set("fs.s3a.proxy.host", "proxy.internal.example.com") + assert( + DeltaScanSupport + .proxyGateReason( + conf, + Seq( + new URI("file:///tmp/table/part-0.parquet"), + new URI("gs://mybucket/part-0.parquet"))) + .isEmpty) + } + + test( + "proxyGateReason declines a plaintext fs.s3a.proxy.host even when a readable global " + + "credential store is configured and clear-text-fallback is false: the host's real " + + "consumer is a plain getTrimmed that consults neither the store nor the fallback flag") { + // getPassword would hide this plaintext host (no store entry, conf fallback disabled), but + // S3AUtils#initProxySupport / AWSClientConfig#createProxyConfiguration read it via plain + // getTrimmed and route Hadoop through the proxy anyway, so the gate must still decline. + withJceks(Map.empty) { conf => + conf.set("hadoop.security.credential.clear-text-fallback", "false") + conf.set("fs.s3a.proxy.host", "proxy.internal.example.com") + val reason = + DeltaScanSupport.proxyGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.proxy.host")) + assert(reason.get.contains("mybucket")) + assert(!reason.get.contains("proxy.internal.example.com")) + } + } + + test( + "proxyGateReason passes when fs.s3a.proxy.host exists only as a global credential-store " + + "entry: the host's real consumer never calls getPassword, so a store-held host cannot " + + "put a proxy into effect") { + // The store entry is real and readable; only lookupPassword-family reads (proxy.username, + // proxy.password) would find it. The host stays empty under plain getTrimmed, so Hadoop + // itself never uses a proxy here and declining would be pure over-refusal. + withJceks(Map("fs.s3a.proxy.host" -> "proxy.internal.example.com")) { conf => + assert( + DeltaScanSupport + .proxyGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } + } + + test( + "proxyGateReason passes on an S3A-scoped provider path when no fs.s3a.proxy.host is set " + + "in plain conf: no keystore, S3A-scoped or otherwise, can supply the host to its real " + + "consumer, so provider configuration alone proves nothing about the proxy") { + // The path points at a nonexistent store on purpose: passing here also proves the gate + // performs no keystore I/O at all for the host, not even to rule the store out. + val tempDir = Files.createTempDirectory("comet-delta-proxy-no-keystore") + try { + val conf = new Configuration(false) + val nonexistentPath = "jceks://file" + tempDir + "/does-not-exist.jceks" + conf.set("fs.s3a.security.credential.provider.path", nonexistentPath) + assert( + DeltaScanSupport + .proxyGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + .isEmpty) + } finally { + Files.delete(tempDir) + } + } + + test( + "proxyGateReason still declines a plaintext fs.s3a.proxy.host when an S3A-scoped provider " + + "path is also configured (plain getTrimmed sees the host regardless of any provider)") { + val tempDir = Files.createTempDirectory("comet-delta-proxy-scoped-provider") + try { + val conf = new Configuration(false) + val nonexistentPath = "jceks://file" + tempDir + "/does-not-exist.jceks" + conf.set("fs.s3a.security.credential.provider.path", nonexistentPath) + conf.set("fs.s3a.proxy.host", "proxy.internal.example.com") + val reason = + DeltaScanSupport.proxyGateReason(conf, Seq(new URI("s3a://mybucket/part-0.parquet"))) + assert(reason.isDefined) + assert(reason.get.contains("fs.s3a.proxy.host")) + assert(reason.get.contains("mybucket")) + } finally { + Files.delete(tempDir) + } + } + + // --------------------------------------------------------------------------------------- + // Discovery harness: mechanically bounds which fs.s3a.* keys this comparator must know about. + // A key cannot compile into the comparator without a consumer tier (AllS3ConfigKeys is + // derived from S3ConfigKeyConsumers), and the tier expectations below pin the assignments. + // (a) AllS3ConfigKeys must cover native's NATIVE_S3A_CONFIG_PROPERTIES + // (native/core/src/parquet/objectstore/s3.rs, kept equal to that file's call sites by a + // Rust test). + // (b) A Rust test in s3.rs checks (a) on every native PR by parsing S3ConfigKeyConsumers + // as text. The parse is mirrored here and must equal AllS3ConfigKeys, so a reformat + // of that block fails this suite rather than silently breaking the Rust check. + // (c) Every credential- or encryption-shaped fs.s3a.* key in Hadoop's Constants class must + // be covered by AllS3ConfigKeys or individually documented as exempt. + // --------------------------------------------------------------------------------------- + + test( + "discovery harness: AllS3ConfigKeys is a superset of native's checked-in " + + "NATIVE_S3A_CONFIG_PROPERTIES list (native/core/src/parquet/objectstore/s3.rs)") { + val rustPath = + DeltaScanContribSuite.findRepoFile("native/core/src/parquet/objectstore/s3.rs") + rustPath match { + case None => + cancel( + "Could not locate native/core/src/parquet/objectstore/s3.rs from this checkout; " + + "skipping the native-key-list superset guard.") + case Some(file) => + val contents = scala.io.Source.fromFile(file, "UTF-8").mkString + val marker = "NATIVE_S3A_CONFIG_PROPERTIES: &[&str] = &[" + val start = contents.indexOf(marker) + assert( + start >= 0, + s"Expected ${file.getAbsolutePath} to declare NATIVE_S3A_CONFIG_PROPERTIES -- has " + + "the constant been renamed or removed?") + val end = contents.indexOf("];", start) + assert(end > start, "Expected a `];`-terminated array literal after the marker") + val arrayBody = contents.substring(start + marker.length, end) + val nativeProperties = + "\"([^\"]*)\"".r.findAllMatchIn(arrayBody).map(_.group(1)).toSet + assert( + nativeProperties.nonEmpty, + "Parsed zero property names out of NATIVE_S3A_CONFIG_PROPERTIES -- the parser above " + + "is likely out of sync with the constant's declaration syntax") + + val nativeKeys = nativeProperties.map(p => s"fs.s3a.$p") + val comparatorKeys = DeltaScanSupport.AllS3ConfigKeys.toSet + val uncovered = nativeKeys.diff(comparatorKeys) + assert( + uncovered.isEmpty, + "Native reads fs.s3a.* key(s) that DeltaScanSupport.AllS3ConfigKeys does not compare, " + + "so a Hadoop-vs-native divergence on any of them would go undetected: " + + s"${uncovered.toSeq.sorted.mkString(", ")} -- add the missing key(s) to " + + "AllS3ConfigKeys") + } + } + + test( + "discovery harness: the text parse of S3ConfigKeyConsumers used by the Rust test in " + + "s3.rs yields exactly AllS3ConfigKeys") { + val scalaPath = "contrib/delta-spark/src/main/scala/org/apache/comet/contrib/delta/" + + "DeltaScanSupport.scala" + DeltaScanContribSuite.findRepoFile(scalaPath) match { + case None => + cancel(s"Could not locate $scalaPath from this checkout; skipping the parse check.") + case Some(file) => + val contents = scala.io.Source.fromFile(file, "UTF-8").mkString + // Same marker and extraction rule as delta_contrib_compares_every_native_s3a_property. + val marker = "S3ConfigKeyConsumers: Seq[(String, S3ConfigConsumer)] = Seq(" + val start = contents.indexOf(marker) + assert( + start >= 0, + s"${file.getAbsolutePath} no longer declares `$marker`: update the parser here and " + + "in s3.rs") + val end = contents.indexOf(")", start + marker.length) + assert(end > start, "Expected a `)` closing the S3ConfigKeyConsumers Seq(") + val block = contents.substring(start + marker.length, end) + val parsed = + "\"(fs\\.s3a\\.[^\"]*)\"\\s*->".r.findAllMatchIn(block).map(_.group(1)).toSet + assert( + parsed == DeltaScanSupport.AllS3ConfigKeys.toSet, + "The Rust test in s3.rs would parse S3ConfigKeyConsumers as " + + s"${parsed.toSeq.sorted.mkString(", ")}, not AllS3ConfigKeys: keep each entry a " + + "single \"fs.s3a.\" -> literal pair, or update both parsers") + } + } + + test( + "discovery harness: every compared key carries exactly one consumer-tier assignment, and " + + "the lookupPassword tier is exactly the credential trio (every other compared key's real " + + "hadoop-aws 3.3.4 consumer is propagateBucketOptions plus a plain Configuration#get " + + "family call, verified per key in S3ConfigKeyConsumers' doc)") { + val keys = DeltaScanSupport.S3ConfigKeyConsumers.map(_._1) + assert( + keys.distinct == keys, + "S3ConfigKeyConsumers assigns more than one tier to the same key -- exactly one " + + "classification per key, declared beside it, is the whole point of the list") + val passwordTier = DeltaScanSupport.S3ConfigKeyConsumers.collect { + case (key, DeltaScanSupport.LookupPasswordConsumer) => key + } + assert( + passwordTier == Seq("fs.s3a.access.key", "fs.s3a.secret.key", "fs.s3a.session.token"), + "The lookupPassword tier changed. A key belongs there ONLY when its real hadoop-aws " + + "consumer is S3AUtils#lookupPassword/#lookupBucketSecret -- verify against the " + + "decompiled call site before updating this expectation, because the wrong tier is not " + + "merely over-cautious: an equality comparator reading wider than the real consumer can " + + "produce a false EQUALITY that admits a diverging scan") + } + + test( + "discovery harness: every credential/encryption-shaped fs.s3a.* key Hadoop's Constants " + + "class declares is either compared by AllS3ConfigKeys or individually documented as " + + "exempt") { + val constantsClassName = "org.apache.hadoop.fs.s3a.Constants" + val constantsClass: Option[Class[_]] = + try { + Some(Class.forName(constantsClassName)) + } catch { + case _: ClassNotFoundException => None + } + constantsClass match { + case None => + cancel( + s"$constantsClassName is not on the test classpath (expected via the " + + "spark-hadoop-cloud test dependency); skipping the sensitive-key coverage guard.") + case Some(cls) => + val allS3aKeys = cls.getFields + .filter { f => + f.getType == classOf[String] && + java.lang.reflect.Modifier.isStatic(f.getModifiers) + } + .flatMap { f => + f.get(null) match { + case s: String if s.startsWith("fs.s3a.") => Some(s) + case _ => None + } + } + .toSet + assert( + allS3aKeys.size > 20, + s"Expected many fs.s3a.* keys via reflection on $constantsClassName, found only " + + s"${allS3aKeys.size} -- has the class's field layout changed in a way this " + + "reflection no longer handles?") + + val sensitiveNameFragments = + Seq("key", "secret", "token", "password", "encryption") + val sensitiveKeys = allS3aKeys.filter { key => + val lower = key.toLowerCase(Locale.ROOT) + sensitiveNameFragments.exists(lower.contains) + } + + val comparatorKeys = DeltaScanSupport.AllS3ConfigKeys.toSet + // Individually justified, one at a time -- NOT a blanket "everything encryption-shaped + // is exempt" carve-out, which would have hidden the SSE-C long-bucket-alias gap just as easily as + // never checking at all. + val documentedExempt: Map[String, String] = Map( + "fs.s3a.encryption.algorithm" -> + ("handled by the dedicated unsupportedEncryptionAlgorithmReason/" + + "effectiveEncryptionAlgorithm allowlist gate, not the generic comparator (needs " + + "its own canonical/deprecated resolution cascade, not a flat single-key compare)"), + "fs.s3a.server-side-encryption-algorithm" -> + "deprecated alias of fs.s3a.encryption.algorithm, same dedicated gate", + "fs.s3a.encryption.key" -> + ("key MATERIAL for the algorithm above; never read for comparison at all -- the " + + "allowlist gate declines on the ALGORITHM alone, so the key's value cannot " + + "change the outcome, and never appears in a decline reason (see " + + "effectiveEncryptionAlgorithm's doc)"), + "fs.s3a.server-side-encryption.key" -> + "deprecated alias of fs.s3a.encryption.key, same reasoning", + "fs.s3a.encryption.cse.kms.region" -> + ("CSE tuning, newer Hadoop only: consulted solely when the algorithm resolves to " + + "a CSE variant, and the allowlist gate declines every CSE algorithm outright, " + + "so this value can never influence an admitted scan; native never reads it"), + "fs.s3a.encryption.cse.custom.keyring.class.name" -> + "CSE tuning, newer Hadoop only, same reasoning as fs.s3a.encryption.cse.kms.region", + "fs.s3a.encryption.cse.v1.compatibility.enabled" -> + "CSE tuning, newer Hadoop only, same reasoning as fs.s3a.encryption.cse.kms.region", + "fs.s3a.proxy.password" -> + ("covered via the dedicated fs.s3a.proxy.host gate (proxyGateReason/" + + "unsupportedProxyReason), not the generic comparator: the password (and the " + + "sibling fs.s3a.proxy.username, not sensitive-shaped so never reaches this map) " + + "only matters once a proxy is actually in effect, and any bucket with a " + + "non-empty effective fs.s3a.proxy.host now declines outright, before any " + + "credential comparison would even run -- so the deployment shape this key used " + + "to be a KNOWN GAP for (a Hadoop deployment requiring a proxy for S3 egress " + + "being silently claimed and connected to directly) can no longer reach this key " + + "at all; the password's VALUE itself is still never read or forwarded to native, " + + "same as before"), + "fs.s3a.failinject.inconsistency.key.substring" -> + ("hadoop-aws test-only S3 fault-injection knob (InconsistentAmazonS3Client " + + "family), not a credential; matches the sensitive-name heuristic only " + + "incidentally via \"key.substring\"")) + + val unclassified = sensitiveKeys + .diff(comparatorKeys) + .diff(documentedExempt.keySet) + assert( + unclassified.isEmpty, + "Hadoop's Constants class declares credential/encryption-shaped fs.s3a.* key(s) " + + "this discovery harness has never classified (neither compared by " + + s"AllS3ConfigKeys nor documented as exempt above): ${unclassified.toSeq.sorted + .mkString(", ")} -- decide whether the key needs a gate, then either add it to " + + "AllS3ConfigKeys or add a justified entry to `documentedExempt` in this test") + } + } + +} + +object DeltaScanContribSuite { + + /** + * Walks up from a candidate root (the `comet.repo.root` system property when set, otherwise + * `user.dir`) looking for `relativePath`. Handles both a repo-root working directory and a + * module-root working directory (e.g. `contrib/delta-spark`) without hardcoding either. + * + * Package-visible (not `private`) so other suites in this package needing a repo-relative file + * (e.g. [[JvmLowercaseParitySuite]]) can share it instead of duplicating it. + */ + private[delta] def findRepoFile(relativePath: String): Option[File] = { + val startDir = Option(System.getProperty("comet.repo.root")) + .map(new File(_)) + .getOrElse(new File(System.getProperty("user.dir"))) + Iterator + .iterate(Option(startDir))(_.flatMap(d => Option(d.getParentFile))) + .takeWhile(_.isDefined) + .map(_.get) + .map(new File(_, relativePath)) + .find(_.isFile) + } +} + +/** A filesystem that records any attempt to open it, then refuses. */ +class OpenRecordingFileSystem extends RawLocalFileSystem { + override def initialize(name: URI, conf: Configuration): Unit = { + OpenRecordingFileSystem.opened.set(true) + throw new IOException(s"$name must not be opened") + } +} + +object OpenRecordingFileSystem { + val opened = new AtomicBoolean(false) +} diff --git a/contrib/delta-spark/src/test/scala/org/apache/spark/sql/comet/DeltaPlanDataInjectorSuite.scala b/contrib/delta-spark/src/test/scala/org/apache/spark/sql/comet/DeltaPlanDataInjectorSuite.scala new file mode 100644 index 00000000000..728dd15cc21 --- /dev/null +++ b/contrib/delta-spark/src/test/scala/org/apache/spark/sql/comet/DeltaPlanDataInjectorSuite.scala @@ -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. + */ + +package org.apache.spark.sql.comet + +import java.util.concurrent.ConcurrentHashMap + +import scala.jdk.CollectionConverters._ + +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.comet.contrib.delta.DeltaSparkScanEnvelope +import org.apache.comet.serde.OperatorOuterClass +import org.apache.comet.serde.OperatorOuterClass.Operator + +/** + * Pins the Delta injector to core's split prepare/inject contract: the partition-invariant common + * is parsed once by [[DeltaPlanDataInjector.prepareCommon]] and shared across tasks through + * core's memo, and [[DeltaPlanDataInjector.inject]] only merges a partition's file list into it. + */ +class DeltaPlanDataInjectorSuite extends AnyFunSuite { + + private val injector = new DeltaPlanDataInjector + + private def commonScan(sourceKey: String): OperatorOuterClass.DeltaSparkScan = + OperatorOuterClass.DeltaSparkScan + .newBuilder() + .setCommon(OperatorOuterClass.NativeScanCommon.newBuilder().setSource("delta-source")) + .setDeltaCommon( + OperatorOuterClass.DeltaSparkScanCommon + .newBuilder() + .setTableRoot("file:/tmp/table") + .setColumnMappingMode("name") + .setSourceKey(sourceKey)) + .build() + + private def partitionScan(paths: String*): OperatorOuterClass.DeltaSparkScan = { + val partition = OperatorOuterClass.DeltaSparkFilePartition.newBuilder() + paths.foreach { path => + partition.addPartitionedFile( + OperatorOuterClass.DeltaSparkPartitionedFile + .newBuilder() + .setFile(OperatorOuterClass.SparkPartitionedFile.newBuilder().setFilePath(path))) + } + OperatorOuterClass.DeltaSparkScan.newBuilder().setFilePartition(partition).build() + } + + private def scanOp(scan: OperatorOuterClass.DeltaSparkScan, children: Operator*): Operator = { + val builder = Operator.newBuilder().setContribScan(DeltaSparkScanEnvelope.pack(scan)) + children.foreach(builder.addChildren) + builder.build() + } + + private def filePaths(op: Operator): Seq[String] = + DeltaSparkScanEnvelope + .unpack(op) + .getFilePartition + .getPartitionedFileList + .asScala + .map(_.getFile.getFilePath) + .toSeq + + test("prepareCommon parses the common half and inject merges only the partition") { + val common = commonScan("delta_k") + val prepared = injector.prepareCommon(common.toByteArray) + assert(prepared == common) + + val op = scanOp(common) + assert(injector.canInject(op)) + assert(injector.getKey(op).contains("delta_k")) + + val injected = + injector.inject(op, prepared, partitionScan("a.parquet", "b.parquet").toByteArray) + val scan = DeltaSparkScanEnvelope.unpack(injected) + assert(scan.getCommon == common.getCommon) + assert(scan.getDeltaCommon == common.getDeltaCommon) + assert(filePaths(injected) == Seq("a.parquet", "b.parquet")) + // A fully populated scan is never a candidate for a second injection. + assert(!injector.canInject(injected)) + } + + test("inject leaves the child list untouched so core can walk it") { + val child = Operator.newBuilder().setPlanId(7).build() + val op = scanOp(commonScan("delta_k"), child) + + val injected = injector.inject( + op, + injector.prepareCommon(commonScan("delta_k").toByteArray), + partitionScan("a.parquet").toByteArray) + + assert(injected.getChildrenCount == 1) + assert(injected.getChildren(0) eq child) + } + + test("core's memo prepares the common once and serves every partition from it") { + val common = commonScan("delta_k").toByteArray + val memo = new ConcurrentHashMap[String, PlanDataInjector.PreparedCommon]() + + val first = PlanDataInjector.prepareShared(injector, "delta_k", common, memo) + val second = PlanDataInjector.prepareShared(injector, "delta_k", common, memo) + assert(second eq first, "a repeat lookup must reuse the parsed common") + assert(memo.size == 1) + + val op = scanOp(commonScan("delta_k")) + val p0 = injector.inject(op, first, partitionScan("p0.parquet").toByteArray) + val p1 = injector.inject(op, second, partitionScan("p1.parquet").toByteArray) + assert(filePaths(p0) == Seq("p0.parquet")) + assert(filePaths(p1) == Seq("p1.parquet")) + } + + test("core's memo replaces a prepared common whose finalized bytes changed under the key") { + val memo = new ConcurrentHashMap[String, PlanDataInjector.PreparedCommon]() + val stale = + PlanDataInjector.prepareShared(injector, "delta_k", commonScan("delta_k").toByteArray, memo) + + val changed = commonScan("delta_k").toBuilder + .setDeltaCommon(commonScan("delta_k").getDeltaCommon.toBuilder.setColumnMappingMode("id")) + .build() + val fresh = PlanDataInjector.prepareShared(injector, "delta_k", changed.toByteArray, memo) + + assert(fresh ne stale) + assert(fresh.getDeltaCommon.getColumnMappingMode == "id") + assert(memo.size == 1, "the stale slot is replaced, not accumulated") + } +} diff --git a/dev/ci/check-ci-config.py b/dev/ci/check-ci-config.py index 3832fda92dc..c752a50c266 100644 --- a/dev/ci/check-ci-config.py +++ b/dev/ci/check-ci-config.py @@ -166,6 +166,21 @@ # neither does the Delta gate, which only inspects build output. (["dev/verify-contrib-delta-gate.sh"], {"delta_gate"}), (["contrib/delta/native/src/lib.rs"], {"delta_gate"}), + # The Delta contrib feeds its own suites and the Linux build, whose strict + # warnings job compiles it on every pull request. A README edit feeds + # neither. + ( + ["contrib/delta-spark/src/main/scala/org/apache/comet/contrib/delta/X.scala"], + { + "build_linux", + "build_linux_full", + "build_linux_all_profiles", + "delta_3_5", + "delta_4_0", + "delta_4_1", + }, + ), + (["contrib/delta-spark/README.md"], set()), ( ["spark/src/test/resources/pyspark/test_pyarrow_udf.py"], { @@ -198,9 +213,10 @@ "delta_gate", "pyarrow_udf", } -# Every other Spark and Iceberg version, and the linux-test matrix's other -# Spark profiles (`build_linux_all_profiles`, part of the Linux build's call -# rather than a job of its own), run once a night against main instead. +# Every other Spark and Iceberg version, every Delta version, and the +# linux-test matrix's other Spark profiles (`build_linux_all_profiles`, part of +# the Linux build's call rather than a job of its own), run once a night +# against main instead. NIGHTLY_TIER = { "spark_3_5", "spark_4_0", @@ -208,6 +224,9 @@ "iceberg_1_8", "iceberg_1_9", "iceberg_1_10", + "delta_3_5", + "delta_4_0", + "delta_4_1", "build_linux_all_profiles", } # Spark 3.4 is deprecated and sits outside the queue and nightly tiers @@ -219,6 +238,8 @@ # One label opts a pull request into every Iceberg version, whichever tier # each sits in. ICEBERG_OPT_IN = {"iceberg_1_8", "iceberg_1_9", "iceberg_1_10", "iceberg_1_11"} +# Likewise one label for every Delta profile. +DELTA_OPT_IN = {"delta_3_5", "delta_4_0", "delta_4_1"} ALL_JOBS = QUEUE_TIER | NIGHTLY_TIER | SPARK_DEPRECATED | {"docs"} assert not QUEUE_TIER & NIGHTLY_TIER, "a job is queue or nightly, never both" @@ -250,6 +271,20 @@ {"name": "pull_request", "action": "synchronize", "labels": ["run-spark-3.5-tests"]}, PR_TIER | {"spark_3_5"}, ), + # The Delta label adds every Delta profile. + ( + {"name": "pull_request", "action": "synchronize", "labels": ["run-delta-tests"]}, + PR_TIER | DELTA_OPT_IN, + ), + ( + { + "name": "pull_request", + "action": "labeled", + "label": "run-delta-tests", + "labels": ["run-delta-tests"], + }, + DELTA_OPT_IN, + ), # The macOS build and the benchmark compile check are queue-only, each # with its own label. Neither label pulls in the other. ( diff --git a/dev/ci/check-suites.py b/dev/ci/check-suites.py index 3acadfb74ad..5b6e0f67e3b 100644 --- a/dev/ci/check-suites.py +++ b/dev/ci/check-suites.py @@ -50,6 +50,10 @@ def file_to_class_name(path: Path) -> str | None: root = Path(".") for path in root.rglob("*Suite.scala"): + # contrib suites run via their own module-level workflows + # (e.g. delta_contrib_test.yml), not the main PR build matrix + if path.parts[0] == "contrib": + continue class_name = file_to_class_name(path) if class_name: if "Shim" in class_name: diff --git a/dev/ci/compute-changes.py b/dev/ci/compute-changes.py index 4d62020bb83..c3786d579c6 100644 --- a/dev/ci/compute-changes.py +++ b/dev/ci/compute-changes.py @@ -44,6 +44,8 @@ "common/**", "spark/**", "spark-integration/**", + # The strict warnings job compiles the Delta contrib against Spark 3.5. + "contrib/delta-spark/**", "pom.xml", "**/pom.xml", ".mvn/**", @@ -435,11 +437,38 @@ ".mvn/**", "mvnw", ], + "delta_3_5": [ + "contrib/delta-spark/**", + "!**.md", + "native/**/src/**", + "native/**/Cargo.toml", + "native/Cargo.lock", + "common/src/main/**", + "common/pom.xml", + "spark/src/main/**", + "!spark/src/main/scala/org/apache/comet/GenerateDocs.scala", + # The contrib suites extend these through the spark test-jar. + "spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala", + "spark/src/test/scala/org/apache/comet/CometS3TestBase.scala", + "spark/pom.xml", + "pom.xml", + "rust-toolchain.toml", + ".github/workflows/ci.yml", + ".github/workflows/delta_contrib_test.yml", + ".github/actions/setup-builder/**", + ], + # Same inputs as delta_3_5: the Delta contrib workflow is called once per + # Spark profile, and the three calls share one path filter. Populated + # below, after the dict, so the lists cannot drift. + "delta_4_0": [], + "delta_4_1": [], } FILTERS["spark_4_1_hive"] = FILTERS["spark_4_1"] FILTERS["build_linux_full"] = FILTERS["build_linux"] FILTERS["build_macos_full"] = FILTERS["build_macos"] FILTERS["build_linux_all_profiles"] = FILTERS["build_linux"] +FILTERS["delta_4_0"] = FILTERS["delta_3_5"] +FILTERS["delta_4_1"] = FILTERS["delta_3_5"] # Which events may run each job, independent of the path filters above. # @@ -577,6 +606,13 @@ "iceberg_1_9": ["nightly", "label:run-iceberg-tests"], "iceberg_1_10": ["nightly", "label:run-iceberg-tests"], "iceberg_1_11": ["queue", "label:run-iceberg-tests"], + # The Delta contrib suites all run nightly. The PR tier's strict warnings + # job compiles the contrib against Spark 3.5, so a change that breaks its + # build on Spark 3.5 fails on its own pull request. One label opts a pull + # request into all three. + "delta_3_5": ["nightly", "label:run-delta-tests"], + "delta_4_0": ["nightly", "label:run-delta-tests"], + "delta_4_1": ["nightly", "label:run-delta-tests"], } diff --git a/dev/verify-contrib-delta-gate.sh b/dev/verify-contrib-delta-gate.sh index 62b0fe4a260..f2ce010e3c3 100755 --- a/dev/verify-contrib-delta-gate.sh +++ b/dev/verify-contrib-delta-gate.sh @@ -17,16 +17,24 @@ # specific language governing permissions and limitations # under the License. # -# Verify the `contrib-delta` build gate keeps Delta surface out of default builds. +# Verify the split between the two Delta features of the native library and the JVM build: +# the kernel-backed `comet-contrib-delta` crate (Cargo feature `contrib-delta`, Maven profile +# `contrib-delta`) stays out of every shipped build, while the small default-on `delta` Cargo +# feature (deletion-vector decoding for the JVM-planned scan) is deliberately in. # # Three independent layers are checked: -# 1. Cargo: default `cargo build` doesn't compile `comet-contrib-delta` and -# doesn't pull `delta_kernel` into the dependency tree. +# 1. Cargo: the default feature set, which is the tree every shipped build compiles, pulls +# neither `comet-contrib-delta` nor `delta_kernel`; the same holds with +# `--no-default-features`. # 2. Maven: default `mvn ... package` doesn't compile any # `org/apache/comet/contrib/` classes and doesn't pull `io.delta:*` deps. -# 3. Symbols: the resulting `libcomet` (`.so` on Linux, `.dylib` on macOS) from the default -# build carries no `comet_contrib_delta`/`delta_kernel`/etc. symbols, and the -# contrib-enabled build carries some (so the pattern is known to still match). +# 3. Symbols: the default `libcomet` (`.so` on Linux, `.dylib` on macOS) carries no +# `comet_contrib_delta`/`delta_kernel`/etc. symbols and does carry the default-on +# deletion-vector decoder, whose symbol footprint is pinned so the feature cannot quietly +# grow into a kernel dependency; the contrib-enabled build carries the contrib symbols. +# File sizes are reported for information only: on an unstripped debug library the +# contrib code is far smaller than build-to-build layout noise, so a size comparison +# cannot tell the two apart. # # Exit non-zero on the first failure. Designed to be wired into CI so a future # change that leaks Delta into core gets caught immediately. @@ -73,25 +81,34 @@ hdr() { printf '\n\033[36m==> %s\033[0m\n' "$*"; } # ---- Cargo gate ----------------------------------------------------------- -hdr "Cargo: default build does not depend on comet-contrib-delta / delta_kernel" +hdr "Cargo: no shipped feature set depends on comet-contrib-delta / delta_kernel" cd "$NATIVE_DIR" -TREE_DEFAULT="$(cargo tree -p datafusion-comet --no-default-features 2>/dev/null)" # Anti-vacuous (mirrors the Maven gate below): a failing `cargo tree` yields empty output, and the # command-substitution failure doesn't trip `set -e` in an assignment -- so assert the root crate we # KNOW is always present before concluding "no Delta deps", otherwise a broken cargo-tree run would # pass the leak check vacuously. (`datafusion-comet ` with a trailing space matches only the root # crate line, not `datafusion-comet-proto`/`-common`.) -if ! grep -q 'datafusion-comet ' <<<"$TREE_DEFAULT"; then - red "FAIL: default cargo tree produced no datafusion-comet entry (cargo tree likely failed;" - red " refusing to conclude 'no Delta deps' vacuously)" - exit 1 -fi -if grep -qE 'comet-contrib-delta|delta_kernel|delta-kernel' <<<"$TREE_DEFAULT"; then - red "FAIL: default cargo tree contains Delta-related deps:" - grep -E 'comet-contrib-delta|delta_kernel|delta-kernel' <<<"$TREE_DEFAULT" - exit 1 -fi -green "OK: cargo tree default is clean of contrib + kernel" +check_tree_clean() { # args: label, then extra `cargo tree` flags + local label="$1" + shift + local tree + tree="$(cargo tree -p datafusion-comet "$@" 2>/dev/null)" + if ! grep -q 'datafusion-comet ' <<<"$tree"; then + red "FAIL: $label cargo tree produced no datafusion-comet entry (cargo tree likely failed;" + red " refusing to conclude 'no Delta deps' vacuously)" + exit 1 + fi + if grep -qE 'comet-contrib-delta|delta_kernel|delta-kernel' <<<"$tree"; then + red "FAIL: $label cargo tree contains Delta-related deps:" + grep -E 'comet-contrib-delta|delta_kernel|delta-kernel' <<<"$tree" + exit 1 + fi + green "OK: $label cargo tree is clean of contrib + kernel" +} +# The default feature set is what every shipped build compiles and includes the `delta` +# feature; `--no-default-features` is the slim opt-out and must stay clean too. +check_tree_clean "default-features" +check_tree_clean "--no-default-features" --no-default-features TREE_CONTRIB="$(cargo tree -p datafusion-comet --features contrib-delta 2>/dev/null)" # The build-gate unit ships a STUB contrib crate, so the gated tree pulls in @@ -227,7 +244,7 @@ green "OK: default build registers no contrib services (empty ServiceLoader regi # ---- libcomet symbol gate ------------------------------------------------- -hdr "libcomet: default build has no Delta symbols" +hdr "libcomet: default build has the delta feature and no contrib symbols, contrib build has them" cd "$NATIVE_DIR" # The cdylib extension is platform-specific: `libcomet.so` on Linux (CI), `libcomet.dylib` on # macOS. Find whichever the build produced; `stat`/`nm` flags also differ across the two. @@ -250,6 +267,27 @@ lib_size() { stat -c%s "$1" 2>/dev/null || stat -f%z "$1"; } delta_syms() { nm "$1" 2>/dev/null | grep -ciE 'comet_contrib_delta|delta_kernel|deltadvfilter|deltasynthetic' || true } +# Symbols of the default-on `delta` feature: the deletion-vector decoder and the planner arms +# that use it. These are what the default build is meant to carry. +DELTA_FEATURE_PATTERN='delta_dv|delta_scan|delta_spark_scan' +delta_feature_syms() { + nm "$1" 2>/dev/null | grep -ciE "$DELTA_FEATURE_PATTERN" || true +} +# Total size in bytes of the default-on `delta` feature's symbols. GNU nm reports sizes with +# `-S`; Mach-O nm always reports zero, so on macOS this returns 0 and the pin below is skipped. +delta_feature_bytes() { + local total=0 size rest + while read -r _ size rest; do + # Undefined symbols carry no size column; skip anything that is not a hex size. + [[ "$size" =~ ^[0-9a-fA-F]+$ && -n "$rest" ]] || continue + total=$((total + 16#$size)) + done < <(nm -S "$1" 2>/dev/null | grep -iE "$DELTA_FEATURE_PATTERN" || true) + echo "$total" +} +# Upper bound for that footprint in an unstripped debug library. It measures 84 KB across 372 +# symbols on Linux; the cap leaves room for toolchain drift but not for a kernel-sized +# dependency riding in through the `delta` feature. +DELTA_FEATURE_MAX_BYTES=$((512 * 1024)) # `nm` is the only direct measurement this section makes, so a missing `nm` has to fail rather # than silently skip -- same anti-vacuous discipline as the cargo-tree and effective-pom guards @@ -272,10 +310,29 @@ fi SIZE_DEFAULT="$(lib_size "$LIB_DEFAULT")" EXT_SYMS="$(delta_syms "$LIB_DEFAULT")" if [[ "$EXT_SYMS" -ne 0 ]]; then - red "FAIL: default libcomet contains $EXT_SYMS Delta-related symbols" + red "FAIL: default libcomet contains $EXT_SYMS contrib/kernel Delta symbols" exit 1 fi -green "OK: default libcomet has 0 Delta symbols (size=$SIZE_DEFAULT bytes)" +green "OK: default libcomet has 0 contrib/kernel symbols (size=$SIZE_DEFAULT bytes)" + +# The default-on `delta` feature must be present and stay small. A default library without +# the decoder means the feature was dropped from the default set; a footprint above the cap +# means something far larger than the decoder now rides in through it. +FEATURE_SYMS="$(delta_feature_syms "$LIB_DEFAULT")" +if [[ "$FEATURE_SYMS" -lt 1 ]]; then + red "FAIL: default libcomet carries no delta feature symbols; the default-on delta feature is missing" + exit 1 +fi +FEATURE_BYTES="$(delta_feature_bytes "$LIB_DEFAULT")" +if [[ "$FEATURE_BYTES" -gt 0 ]]; then + if [[ "$FEATURE_BYTES" -gt "$DELTA_FEATURE_MAX_BYTES" ]]; then + red "FAIL: default-on delta feature symbols total $FEATURE_BYTES bytes, above the $DELTA_FEATURE_MAX_BYTES byte cap" + exit 1 + fi + green "OK: default libcomet carries the delta feature ($FEATURE_SYMS symbols, $FEATURE_BYTES bytes, cap $DELTA_FEATURE_MAX_BYTES)" +else + green "OK: default libcomet carries the delta feature ($FEATURE_SYMS symbols; nm reports no sizes on this platform, footprint cap not checked)" +fi cargo build -j 4 -p datafusion-comet --features contrib-delta >/dev/null 2>&1 LIB_CONTRIB="$(comet_lib)" @@ -310,8 +367,9 @@ green "OK: contrib-enabled libcomet has $CONTRIB_SYMS Delta symbols (size=$SIZE_ # ---- Summary -------------------------------------------------------------- hdr "All gate checks passed" -echo " default cargo: no comet-contrib-delta, no delta_kernel" +echo " default cargo: no comet-contrib-delta, no delta_kernel (default and --no-default-features)" echo " default mvn: no io.delta:*, no contrib/delta classes" -echo " default dylib: 0 Delta symbols (contrib build has $CONTRIB_SYMS)" +echo " default dylib: delta feature present ($FEATURE_SYMS symbols), 0 contrib/kernel symbols" +echo " contrib dylib: $CONTRIB_SYMS contrib/kernel symbols" echo echo "Run with: dev/verify-contrib-delta-gate.sh" diff --git a/docs/source/contributor-guide/ci.md b/docs/source/contributor-guide/ci.md index e3ebece2397..c567e6f0a75 100644 --- a/docs/source/contributor-guide/ci.md +++ b/docs/source/contributor-guide/ci.md @@ -45,7 +45,7 @@ flowchart LR PRTIER["PR tier
Linux build, lint, Rust tests
TPC-H / TPC-DS
Comet suites, Spark 4.1"] QUEUE["Queue tier, on top of the PR tier
Spark SQL, Spark 4.1
Iceberg 1.11
macOS build and Comet suites
Benchmark check, Delta gate
PyArrow UDF, Spark 4.0 / 4.1 / 4.2"] - NIGHTLY["Nightly tier
Comet suites, Spark 3.4 / 3.5 / 4.0 / 4.2
Spark SQL, Spark 3.5 / 4.0 / 4.2
Iceberg 1.8 / 1.9 / 1.10"] + NIGHTLY["Nightly tier
Comet suites, Spark 3.4 / 3.5 / 4.0 / 4.2
Spark SQL, Spark 3.5 / 4.0 / 4.2
Iceberg 1.8 / 1.9 / 1.10
Delta contrib, Spark 3.5 / 4.0 / 4.1"] S34["Neither tier
Spark SQL, Spark 3.4"] CACHE["Cache-refresh-only mode
the three cache-writing jobs, plus Lint"] ``` @@ -66,13 +66,15 @@ Suite by suite: | Comet test suites, Spark 3.4 / 3.5 / 4.0 / 4.2 | with label | no | yes | | Spark SQL tests, Spark 3.5 / 4.0 / 4.2 | with label | no | yes | | Iceberg Spark SQL tests, Iceberg 1.8 / 1.9 / 1.10 | with label | no | yes | +| Delta contrib tests, Spark 3.5 / 4.0 / 4.1 | with label | no | yes | | Spark SQL tests, Spark 3.4 | with label | no | no | The **PR tier** is the fast feedback loop while a change is being iterated on. The **queue tier** is the authoritative gate: everything the PR tier runs plus the suites that must pass before a change lands, evaluated against the merge result rather than the pull request head. The queue runs one Spark version and one Iceberg version, the ones the default build profile targets. The -**nightly tier** runs the other Spark and Iceberg versions once a day against `main` as it stands; +**nightly tier** runs the other Spark and Iceberg versions and every Delta contrib suite once a day +against `main` as it stands; see [Nightly runs](#nightly-runs) below. Nothing in the queue tier runs again on push to `main`, because the queue already tested the exact tree that landed. The exceptions are the Linux and macOS builds, which also run on push so that the dependency caches on `main` stay @@ -120,6 +122,7 @@ Each suite outside the PR tier has a label that runs it on a pull request: | `run-benchmark-check` | Benchmark compile and lint check | | `run-delta-build-gate` | Delta contrib build gate | | `run-pyarrow-udf-tests` | PyArrow UDF tests against Spark 4.0/4.1/4.2 | +| `run-delta-tests` | Delta contrib tests against Spark 3.5, 4.0 and 4.1 | | `run-spark-4.1-tests` | Spark SQL tests against Spark 4.1, every module | | `run-spark-4.1-hive-tests` | Spark SQL tests against Spark 4.1, sql_hive only | | `run-spark-3.4-tests` | Spark SQL tests against Spark 3.4 | @@ -172,6 +175,8 @@ does not cover. Some examples: - anything that touches Hive table support, `InsertIntoHiveTable`, or the `sql/hive` parts of the 4.1 diff - anything touching Iceberg reflection, the Iceberg scan or write path, or the Iceberg diffs +- a change to `contrib/delta-spark/` or to the scan and native code it builds on; + `run-delta-tests` runs the Delta contrib suites against Spark 3.5, 4.0 and 4.1 - native code with platform-specific behavior, or a dependency bump that changes what is compiled on macOS - a change to the benchmark sources under `spark/src/test/scala/org/apache/spark/sql/benchmark` @@ -246,8 +251,9 @@ re-armed automatically. `ci.yml` also runs on a schedule, at 06:00 UTC every day, against what landed on `main` since the last successful scheduled run. That run executes only the nightly tier: the Comet test suites against the Spark -profiles other than 4.1, the Spark SQL suites for Spark 3.5, 4.0 and 4.2, and the Iceberg suites for -1.8, 1.9 and 1.10, each only when the day's changes touched files it covers. The queue already ran +profiles other than 4.1, the Spark SQL suites for Spark 3.5, 4.0 and 4.2, the Iceberg suites for +1.8, 1.9 and 1.10, and the Delta contrib suites for Spark 3.5, 4.0 and 4.1, each only when the day's +changes touched files it covers. The queue already ran everything else against the same tree, so nothing in the queue tier is repeated. A day with no merges, or with only documentation changes, runs nothing. diff --git a/docs/source/contributor-guide/development.md b/docs/source/contributor-guide/development.md index d1190124016..ea4a65e3089 100644 --- a/docs/source/contributor-guide/development.md +++ b/docs/source/contributor-guide/development.md @@ -669,7 +669,7 @@ Make sure to resolve any Clippy warnings before submitting your pull request, as The `Strict Scala warnings` job runs on every pull request and in the merge queue. It compiles the main and test sources with scalac warnings promoted to errors, so anything it reports fails the build — an `Int` widened into a `Long` metric, or a discarded builder result, for example. Reproduce it locally with: ```sh -./mvnw test-compile -Pspark-3.5 -Pstrict-warnings -DskipTests +./mvnw test-compile -Pspark-3.5,delta -Pstrict-warnings -DskipTests ``` Use the Spark 3.5 profile: it is the one the job runs, and the default build profile will not reproduce it. The default is Spark 4.1 on Scala 2.13, where `-Pstrict-warnings` still fails on warnings unrelated to your change (tracked in [#5893](https://github.com/apache/datafusion-comet/issues/5893)), and where the compiler reports a different set — an adapted argument list, for instance, is flagged under Scala 2.12 but not under 2.13. diff --git a/docs/source/user-guide/latest/delta.md b/docs/source/user-guide/latest/delta.md new file mode 100644 index 00000000000..2d8d7872b30 --- /dev/null +++ b/docs/source/user-guide/latest/delta.md @@ -0,0 +1,74 @@ + + +# Delta Lake (experimental) + +Comet can execute DSv1 Delta Lake table scans natively. Reads planned by +delta-spark run through Comet's native Parquet scan, inheriting row-group +pruning, page-index pruning, and filter pushdown, with deletion vectors +applied inside the scan. + +Support is experimental and explicitly opt-in. Two things are required: + +1. The `comet-contrib-delta-spark` contrib jar on the classpath, alongside + `delta-spark`. It is never bundled into `comet-spark`, and no release + publishes it yet, so build it from source with the steps in the + [contrib module README](https://github.com/apache/datafusion-comet/blob/main/contrib/delta-spark/README.md#building-and-testing). +2. `spark.comet.scan.delta.enabled=true`. The default is `false`, so + the jar alone does nothing. + +Unsupported tables and features fall back to Spark's reader. See the +[contrib module README](https://github.com/apache/datafusion-comet/blob/main/contrib/delta-spark/README.md) +for the supported Spark/Delta version matrix and build instructions. + +Unlike the core native scan, the Delta scan resolves each data file's datetime +calendar-rebase policy from the file's own writer metadata +(`org.apache.spark.legacyDateTime` and friends), the same way Spark's reader +does, selecting the `datetimeRebaseModeInRead` spec for dates and INT64 +timestamps and the `int96RebaseModeInRead` spec for INT96 timestamps, at any +nesting depth. As in Spark, the spec follows the type a column is read as: a +timestamp column read as `TIMESTAMP_NTZ` is never rebased (a `DATE` column read +as `TIMESTAMP_NTZ` keeps the date spec), and a column read as `TIMESTAMP` takes +the datetime spec even when the file marks it as not adjusted to UTC. Dates written with the legacy hybrid Julian/Gregorian calendar +are rebased exactly, timestamps are rebased exactly when the file records a +fixed UTC writer time zone, and ancient values whose calendar cannot be +applied natively (non-UTC legacy writer zones, or files that do not declare a +policy under the `EXCEPTION` read mode) raise an error rather than silently +returning shifted values. Modern values are unaffected: dates from 1582-10-15 +onward, and timestamps from 1900-01-01T00:00:00Z onward (Spark's own +rebase cutoff). Disable `spark.comet.scan.delta.enabled` for such tables to +read them through Spark. + +## Configuration + + + +| Config | Description | Default Value | +|--------|-------------|---------------| +| `spark.comet.scan.delta.dv.maxDeletedRowsPerFile` | Upper bound on a single file's deletion-vector cardinality (deleted row count) the native Delta scan will claim. Applying a deletion vector expands it into per-row selectors that are held in memory. This bound caps one file's selectors, not what a task holds: the selectors for every file in a partition stay held until the task finishes. The bound is a deliberately pessimistic planning-time proxy for that memory (deletion vector cardinality, not the exact selector count), so a large but contiguous deletion is declined the same as a large alternating one. Scans whose deletion vectors exceed this bound for any file fall back to Spark's reader. | 1000000 | +| `spark.comet.scan.delta.enabled` | Whether to enable native Delta table scans. When enabled, DSv1 Delta table reads planned by delta-spark are executed through Comet's native Parquet scan, inheriting row-group pruning, page-index pruning, and filter pushdown, with deletion vectors applied inside the scan. Experimental: defaults to false, so adding the contrib jar does not by itself change how any query is read. | false | + + +When Spark splits a Parquet file into several tasks, each native Delta scan builds +deletion-vector row selectors only for the row groups owned by its split. Deletion +positions remain relative to the entire file, and positions beyond the file's row +count are still rejected. Each split still fetches and decodes the full deletion +vector and uses the existing conservative whole-file memory admission bound before +building selectors. The per-file cardinality limit above is unchanged. diff --git a/docs/source/user-guide/latest/index.rst b/docs/source/user-guide/latest/index.rst index 653bdf37889..84b0064647d 100644 --- a/docs/source/user-guide/latest/index.rst +++ b/docs/source/user-guide/latest/index.rst @@ -95,6 +95,7 @@ the navigation menu to read more. :caption: Integrations :hidden: + Delta Lake (experimental) Iceberg Guide Iceberg Writes Celeborn Guide diff --git a/native/Cargo.lock b/native/Cargo.lock index e4d2aac7e1d..7bc355b68eb 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -2038,6 +2038,7 @@ dependencies = [ "reqsign-aws-v4", "reqsign-core", "reqwest 0.12.28", + "roaring", "serde", "serde_json", "tempfile", diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml index b21c5bc75a8..d52c6f96590 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -35,6 +35,9 @@ include = [ publish = false [dependencies] +# Delta deletion-vector decoding (feature = "delta") +roaring = { version = "0.11", optional = true } +crc32fast = { version = "1.5", optional = true } arrow = { workspace = true } base64 = "0.23.0" chrono-tz = { workspace = true } @@ -122,13 +125,25 @@ reqsign-aws-v4 = "3" [features] backtrace = ["datafusion/backtrace"] -default = ["hdfs-opendal"] +default = ["hdfs-opendal", "delta"] contrib-lance = ["dep:comet-contrib-lance"] hdfs-opendal = ["opendal/services-hdfs", "object_store_opendal", "hdfs-sys"] jemalloc = ["tikv-jemallocator", "tikv-jemalloc-ctl"] -# Delta Lake integration. When enabled, links the `comet-contrib-delta` crate -# into `libcomet` and activates the `OpStruct::DeltaScan` dispatcher arm. -# Default builds carry zero Delta surface. +# Native Delta Lake scan support for the JVM-planned path (contrib/delta-spark). +# In the default set so trying the contrib needs only the jar and the config, +# not a custom native build. It is inert at runtime unless the contrib jar is +# on the classpath (ServiceLoader) AND spark.comet.scan.delta.enabled is set, +# so it cannot affect non-Delta scans, and it adds no crate to the default +# build: `roaring` and `crc32fast` are already in the tree through iceberg and +# the shuffle crate, so the feature only promotes two transitive dependencies +# to direct ones. dev/verify-contrib-delta-gate.sh measures and caps its +# footprint. Opt out with --no-default-features for slim builds; the planner +# arm then returns a clear "built without the delta feature" error. +delta = ["dep:roaring", "dep:crc32fast"] +# Delta Lake integration via delta-kernel-rs. When enabled, links the +# `comet-contrib-delta` crate into `libcomet` and activates the contrib scan +# dispatcher arm. Default builds carry zero delta-kernel surface; the `delta` +# feature above has no kernel dependency. contrib-delta = ["dep:comet-contrib-delta"] # exclude optional packages from cargo machete verifications diff --git a/native/core/src/execution/delta_dv.rs b/native/core/src/execution/delta_dv.rs new file mode 100644 index 00000000000..982c4434207 --- /dev/null +++ b/native/core/src/execution/delta_dv.rs @@ -0,0 +1,2748 @@ +// 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. + +//! Delta Lake deletion-vector decoding and translation into DataFusion +//! [`ParquetAccessPlan`]s (feature = "delta"). +//! +//! Wire formats implemented here (from delta-spark's `DeletionVectorStore` / +//! `RoaringBitmapArray`, v3.3.2): +//! - On-disk DV file: 1 version byte at the start of the file; at +//! `descriptor.offset`: `[i32 BE size][data: size bytes][i32 BE CRC32(data)]`. +//! - `data`: `[i32 LE magic]` then either +//! - magic 1681511376 ("native"): `[i32 LE count]`, then per bitmap +//! `[i32 LE size][standard 32-bit RoaringBitmap]`, keys implicit (index); +//! - magic 1681511377 ("portable", the spec's 64-bit extension): `[i64 LE +//! count]`, then per bitmap `[i32 LE key][standard 32-bit RoaringBitmap]` +//! with keys ascending -- exactly [`RoaringTreemap`]'s serialized form. + +use std::mem::size_of; +use std::sync::Arc; + +use datafusion::datasource::listing::PartitionedFile; +use datafusion::datasource::physical_plan::parquet::metadata::DFParquetMetadata; +use datafusion::datasource::physical_plan::parquet::{ParquetAccessPlan, RowGroupAccess}; +use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion::execution::runtime_env::RuntimeEnv; +use datafusion_datasource::FileRange; +use futures::{StreamExt, TryStreamExt}; +use object_store::path::Path; +use object_store::{ObjectStore, ObjectStoreExt}; +use parquet::arrow::arrow_reader::{RowSelection, RowSelector}; +use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData}; +use roaring::{RoaringBitmap, RoaringTreemap}; + +use crate::execution::operators::ExecutionError; +use crate::execution::operators::ExecutionError::GeneralError; +use datafusion_comet_proto::spark_operator::DeltaSparkDvDescriptor; + +const NATIVE_MAGIC: i32 = 1681511376; +const PORTABLE_MAGIC: i32 = 1681511377; + +/// Unframe a DV blob read from `descriptor.offset` of a DV file: +/// `[i32 BE size][data][i32 BE crc]`. Verifies both the size against the +/// descriptor's `size_in_bytes` and the CRC32 checksum. +/// An inline payload carries no framing, so its length is checked against the descriptor here, +/// the way `unframe_dv_blob` checks an on-disk blob's size header. +fn check_inline_payload_size( + file_path: &str, + payload: &[u8], + size_in_bytes: i32, +) -> Result<(), ExecutionError> { + if payload.len() as i64 != i64::from(size_in_bytes) { + return Err(GeneralError(format!( + "Inline deletion vector for {file_path} has {} bytes but its descriptor says {size_in_bytes}", + payload.len() + ))); + } + Ok(()) +} + +pub fn unframe_dv_blob(blob: &[u8], expected_size: usize) -> Result<&[u8], ExecutionError> { + if blob.len() < 8 { + return Err(GeneralError(format!( + "Deletion vector blob too short: {} bytes", + blob.len() + ))); + } + let size = i32::from_be_bytes(blob[0..4].try_into().unwrap()); + if size < 0 || size as usize != expected_size { + return Err(GeneralError(format!( + "Deletion vector size mismatch: file says {size}, descriptor says {expected_size}" + ))); + } + let end = 4 + size as usize; + if blob.len() < end + 4 { + return Err(GeneralError(format!( + "Deletion vector blob truncated: need {} bytes, have {}", + end + 4, + blob.len() + ))); + } + let data = &blob[4..end]; + let expected_crc = i32::from_be_bytes(blob[end..end + 4].try_into().unwrap()); + let actual_crc = crc32fast::hash(data) as i32; + if expected_crc != actual_crc { + return Err(GeneralError( + "Deletion vector checksum mismatch".to_string(), + )); + } + Ok(data) +} + +/// Deserialize the magic-prefixed RoaringBitmapArray into a 64-bit treemap of +/// deleted row indexes. +pub fn deserialize_dv_bitmap(data: &[u8]) -> Result { + if data.len() < 4 { + return Err(GeneralError( + "Deletion vector bitmap too short for magic number".to_string(), + )); + } + let magic = i32::from_le_bytes(data[0..4].try_into().unwrap()); + let rest = &data[4..]; + match magic { + PORTABLE_MAGIC => RoaringTreemap::deserialize_from(rest) + .map_err(|e| GeneralError(format!("Invalid portable deletion vector bitmap: {e}"))), + NATIVE_MAGIC => { + if rest.len() < 4 { + return Err(GeneralError( + "Native deletion vector bitmap missing count".to_string(), + )); + } + let count = i32::from_le_bytes(rest[0..4].try_into().unwrap()); + if count < 0 { + return Err(GeneralError(format!( + "Invalid RoaringBitmapArray length ({count} < 0)" + ))); + } + let mut pos = 4usize; + let mut treemap = RoaringTreemap::new(); + for key in 0..count as u64 { + if rest.len() < pos + 4 { + return Err(GeneralError( + "Native deletion vector bitmap truncated".to_string(), + )); + } + let size = i32::from_le_bytes(rest[pos..pos + 4].try_into().unwrap()); + pos += 4; + if size < 0 || rest.len() < pos + size as usize { + return Err(GeneralError( + "Native deletion vector bitmap truncated".to_string(), + )); + } + let bitmap = RoaringBitmap::deserialize_from(&rest[pos..pos + size as usize]) + .map_err(|e| { + GeneralError(format!("Invalid deletion vector sub-bitmap: {e}")) + })?; + pos += size as usize; + for value in bitmap { + treemap.insert((key << 32) | value as u64); + } + } + Ok(treemap) + } + other => Err(GeneralError(format!( + "Unexpected RoaringBitmapArray magic number {other}" + ))), + } +} + +/// Translate deleted row indexes into a [`ParquetAccessPlan`]: fully-deleted +/// row groups become `Skip`, untouched groups stay `Scan`, and partially +/// deleted groups get a `RowSelection` selecting the complement of the deleted +/// rows. Page-index pruning later INTERSECTS with these selections, so DV +/// skips and page skips compose. When a file range is present, only its row groups acquire +/// selectors; skipped groups still advance the file-global deletion-vector positions. +pub fn build_access_plan( + metadata: &ParquetMetaData, + deleted: &RoaringTreemap, + file_range: Option<&FileRange>, +) -> Result { + let mut plan = ParquetAccessPlan::new_all(metadata.num_row_groups()); + // Empty bitmaps create no selectors to avoid. Let the reader apply its ordinary range + // pruning instead of doing an extra page-offset pass for this case. + let file_range = file_range.filter(|_| !deleted.is_empty()); + let mut deleted_iter = deleted.iter(); + let mut next_deleted = deleted_iter.next(); + let mut group_start = 0u64; + for (idx, row_group) in metadata.row_groups().iter().enumerate() { + let num_rows = row_group.num_rows(); + // A corrupt footer can report a negative row count. `num_rows as u64` would otherwise + // wrap it into a huge positive value, silently corrupting every row-group boundary + // computed from `group_start`/`group_end` below (and therefore which deleted row indexes + // land in which row group) instead of failing loudly. + if num_rows < 0 { + return Err(GeneralError(format!( + "Parquet footer reports a negative row count ({num_rows}) for row group {idx}" + ))); + } + let num_rows = num_rows as u64; + let group_end = group_start.checked_add(num_rows).ok_or_else(|| { + GeneralError(format!( + "Parquet footer row counts overflow at row group {idx} ({group_start} + {num_rows})" + )) + })?; + // Match DataFusion's prune_by_range: dictionary-page start takes precedence over + // data-page start, and FileRange is half-open. Keep an entry for every row group and + // advance file-global row positions even when this split cannot read the group. + if file_range.is_some_and(|range| { + let column = row_group.column(0); + let offset = column + .dictionary_page_offset() + .unwrap_or_else(|| column.data_page_offset()); + !range.contains(offset) + }) { + plan.skip(idx); + group_start = group_end; + continue; + } + + // Seeking skips deletion positions from preceding, excluded groups without visiting + // each one. The pinned Roaring implementation also handles absent 32-bit containers. + if next_deleted.is_some_and(|row| row < group_start) { + deleted_iter.advance_to(group_start); + next_deleted = deleted_iter.next(); + } + let mut selectors: Vec = Vec::new(); + let mut cursor = group_start; + let mut deleted_in_group = 0u64; + while let Some(row) = next_deleted.filter(|row| *row < group_end) { + next_deleted = deleted_iter.next(); + deleted_in_group += 1; + if row > cursor { + selectors.push(RowSelector::select((row - cursor) as usize)); + } + // Merge runs of consecutive deleted rows into one skip. + match selectors.last_mut() { + Some(last) if last.skip => last.row_count += 1, + _ => selectors.push(RowSelector::skip(1)), + } + cursor = row + 1; + } + if deleted_in_group == num_rows && num_rows > 0 { + plan.skip(idx); + } else if deleted_in_group > 0 { + if group_end > cursor { + selectors.push(RowSelector::select((group_end - cursor) as usize)); + } + plan.scan_selection(idx, RowSelection::from(selectors)); + } + group_start = group_end; + } + // Validate against the whole file even when this split excludes the final groups. + // A deleted index beyond the file's total row count means the DV does not belong to this + // file. Checking the maximum also catches corruption beyond an absent Roaring container. + if let Some(row) = deleted.max().filter(|row| *row >= group_start) { + return Err(GeneralError(format!( + "Deletion vector marks row {row} but the file only has {group_start} rows" + ))); + } + Ok(plan) +} + +/// Verify a decoded deletion vector's row count matches the descriptor's +/// declared `cardinality`, mirroring Delta's JVM reader +/// (`StoredBitmap.validateCardinality`). The CRC and framing checks catch +/// corruption but not a stale, otherwise well-formed bitmap whose row count +/// no longer matches the descriptor -- that would silently under- or +/// over-delete rows. +fn validate_cardinality( + file_path: &str, + expected: i64, + deleted: &RoaringTreemap, +) -> Result<(), ExecutionError> { + let actual = deleted.len(); + if actual != expected as u64 { + return Err(GeneralError(format!( + "Deletion vector for {file_path} has cardinality mismatch: descriptor says {expected}, decoded bitmap has {actual} deleted rows" + ))); + } + Ok(()) +} + +/// One data file plus everything needed to apply its deletion vector. The +/// file's size comes from `file.object_meta.size` (built by the planner from +/// the proto's `file_size`). +/// +/// `data_store` and `dv_store` are resolved by the caller *before* entering +/// the async `attach_access_plans` runtime (see its doc comment): building an +/// object store is sync I/O that, for a cold S3 authority, internally issues +/// its own `Handle::block_on` calls, which panics if nested inside another +/// `block_on`. Resolving up front means this module never constructs a +/// store itself. +pub struct DvScanFile { + pub file: PartitionedFile, + /// Full URL of the data file (proto `file_path`). + pub file_path: String, + pub dv: Option, + /// Object store for `file_path`, pre-resolved by the caller. Only read + /// when `dv` is `Some` (files without a deletion vector never open their + /// footer here), but every file carries one so the struct's shape + /// doesn't depend on whether a deletion vector is present. + pub data_store: Arc, + /// Store and within-store path for an on-disk deletion vector's absolute + /// path, pre-resolved by the caller. `None` when the file has no + /// deletion vector or the deletion vector is stored inline. + pub dv_store: Option<(Arc, Path)>, +} + +/// Execution-memory-pool reservation covering one file's expanded DV row selectors across +/// their *entire* lifetime attached to a scan -- from `build_access_plan`'s construction +/// through DataFusion 54.1's reader normalizing the attached [`ParquetAccessPlan`] +/// (`create_initial_plan`'s deep clone plus `into_overall_row_selection`'s combined +/// `RowSelection`; see [`reader_peak_bytes`]) -- attached to the file's [`PartitionedFile`] +/// extensions alongside its [`ParquetAccessPlan`]. The reservation's lifetime is tied to the +/// `PartitionedFile` it is attached to, so it is released back to the pool exactly when the +/// plan is dropped (query completion or an early-terminated scan), never held open longer. +/// Every file in one [`attach_access_plans`] call draws its reservation from the same +/// registered consumer, so the pool counts one consumer for the partition while any of those +/// files is alive, not one per file. Newtype-wrapped so it occupies its own slot in the +/// multi-slot, type-keyed `extensions` map (`datafusion_common::extensions::Extensions`) +/// alongside the plan, rather than a bare `MemoryReservation` colliding with one some other +/// extension might attach. +pub struct DvAccessPlanReservation(pub MemoryReservation); + +/// Total number of [`RowSelector`]s materialized across `plan`'s per-row-group +/// selections (`RowGroupAccess::Selection`); `Scan`/`Skip` row groups +/// contribute none. An alternating deleted/retained bitmap produces one +/// non-coalescing selector per row (see [`reader_peak_bytes`]'s doc comment +/// for the worst-case accounting), so this count -- not the deletion +/// vector's cardinality -- is the thing that must be bounded and reserved +/// against the execution memory pool. +fn total_selectors(plan: &ParquetAccessPlan) -> usize { + plan.inner() + .iter() + .map(|access| match access { + RowGroupAccess::Selection(selection) => selection.iter().count(), + _ => 0, + }) + .sum() +} + +/// Multiplier bounding the peak allocation live *during construction* of one +/// file's [`RowSelection`]s, relative to the conservative selector-count +/// bound `S = 2 * cardinality + num_row_groups` (one non-coalescing selector +/// per deleted row in the worst-case alternating pattern, doubled, plus up to +/// one extra boundary selector per row group). Split `S` into `r`, the +/// selectors already retained from row groups `build_access_plan` has +/// finished, and `c`, the selectors accumulated so far in the current row +/// group's source `Vec`; `r` and `c` partition the selectors counted toward +/// `S`, so `r + c <= S` always. While the current group is being built, the +/// `Vec`'s doubling growth strategy can leave its backing allocation at up to +/// `2 * c` (the next power-of-two capacity above `c`). Once the group +/// finishes, `RowSelection::from(Vec)` (parquet's `FromIterator` impl, +/// `with_capacity` + copy) builds a second, separate `Vec` of size `c` from +/// that source while the source is still alive, so at the moment the copy +/// begins, the retained selectors, the current group's doubled source `Vec`, +/// and the copy are all live simultaneously: `r + 2c + c = r + 3c`. Since +/// `r >= 0`, `r + 3c <= 3r + 3c = 3(r + c) <= 3S`. 3x covers that peak. +const CONSTRUCTION_PEAK_FACTOR: usize = 3; + +/// Upper bound on how much larger a `Vec`'s backing allocation can be than its element count +/// after being built by repeated pushes: `std`'s doubling growth strategy never leaves a `Vec` +/// of `n` elements with a backing allocation larger than the next power of two above `n`, which +/// is at most `2 * n` for any `n >= 1`. +const VEC_GROWTH_CAPACITY_FACTOR: usize = 2; + +/// `RawVec`'s minimum non-zero capacity for element sizes `<= 1024` bytes ([`RowSelector`] is +/// 16 bytes on 64-bit platforms: a `usize` row count plus a padded `bool`). Applied once per +/// row group (or per contiguous run of row groups) a fresh `from_fn`/`FlatMap`-driven `Vec` +/// gets built for (see [`reader_peak_bytes`]), so even a group or run whose true selector count +/// is tiny still pays this floor. +const MIN_VEC_CAPACITY_SELECTORS: usize = 4; + +/// Conservative upper bound, in bytes, on the peak allocation live while DataFusion 54.1's +/// reader normalizes one file's attached [`ParquetAccessPlan`]. This, not the plan's own +/// retained selector bytes, is what the steady-state reservation must cover. Three allocations +/// can be live at once when `into_overall_row_selection` returns: +/// +/// 1. **Attached original** (`selectors`, exact): `create_initial_plan` deep-clones the plan +/// while the original stays reachable from the file's `extensions`. The original's `Vec`s +/// have no excess capacity, because `RowSelection::from(Vec)` (what +/// `build_access_plan` uses) is a `with_capacity(len)` copy. +/// 2. **The clone** (`<= VEC_GROWTH_CAPACITY_FACTOR * selectors + MIN_VEC_CAPACITY_SELECTORS * +/// num_row_groups`): exact unless page-index pruning fires. When it does, +/// `PagePruningAccessPlanFilter` (via `access_plan.rs`'s `scan_selection`) replaces each of +/// the clone's `RowGroupAccess::Selection`s with `RowSelection::intersection`'s output, +/// which `intersect_row_selections` builds with a `from_fn` generator whose `size_hint()` is +/// `(0, None)`. Each row group's `Vec` therefore starts empty and doubles as it grows, and +/// since row groups are intersected independently the capacity floor is paid per row group. +/// 3. **Per-run combined selection** (`<= VEC_GROWTH_CAPACITY_FACTOR * (selectors + +/// num_row_groups) + MIN_VEC_CAPACITY_SELECTORS * num_row_groups`): +/// `into_overall_row_selection` collects each contiguous run of row groups into a new +/// `RowSelection` through a `FlatMap` whose `size_hint().0` is 0, so a run of `len` +/// selectors allocates at most `MIN_VEC_CAPACITY_SELECTORS + VEC_GROWTH_CAPACITY_FACTOR * +/// len`. `len` is the run's share of `selectors` plus one `RowSelector::select(num_rows)` per +/// `RowGroupAccess::Scan` row group, and there are at most `num_row_groups` runs. +/// +/// In bytes, with the constants above: `(5 * selectors + 10 * num_row_groups) * +/// size_of::()`. Checked against two measured worst cases: +/// +/// - No page-index pruning (term 2 stays exact): one 2,000,000-row group, 1,000,000 +/// alternating deletions, `selectors = 2,000,000`. Measured allocator peak 97,554,457 B, of +/// which the three terms account for 97,554,432 B (a 25 B residue is unattributed). This +/// bound gives 160,000,160 B, looser because it must also cover the next case. +/// - Page-index pruning fires against the clone: one 1,048,577-row group, `selectors = +/// 1,048,577`. Measured peak 83,886,096 B; this bound gives 83,886,320 B, a deliberately +/// tight 224 B margin since this case drives the bound. +/// +/// Uses checked arithmetic: a count large enough to overflow `usize` means corrupt or malicious +/// input and is reported as an error rather than a panic. +fn reader_peak_bytes(selectors: usize, num_row_groups: usize) -> Result { + let overflow = || { + GeneralError(format!( + "Deletion vector reader-peak bound overflowed for {selectors} selectors and \ + {num_row_groups} row groups" + )) + }; + // Term 1: the attached original -- exact, untouched by page-index pruning (only the clone + // is ever intersected; see the doc comment above). + let attached_term = selectors; + // Term 2: the clone, bounded as if page-index pruning DID fire against every row group + // (safe even when it doesn't: term 2's bound is always >= `selectors`, so it never + // undershoots the exact case either). + let clone_growth = selectors + .checked_mul(VEC_GROWTH_CAPACITY_FACTOR) + .ok_or_else(overflow)?; + let clone_floor = num_row_groups + .checked_mul(MIN_VEC_CAPACITY_SELECTORS) + .ok_or_else(overflow)?; + let clone_term = clone_growth.checked_add(clone_floor).ok_or_else(overflow)?; + // Term 3: into_overall_row_selection's per-run combined-selection allocation. + let combined_growth = selectors + .checked_mul(VEC_GROWTH_CAPACITY_FACTOR) + .ok_or_else(overflow)?; + let combined_floor = num_row_groups + .checked_mul(MIN_VEC_CAPACITY_SELECTORS + VEC_GROWTH_CAPACITY_FACTOR) + .ok_or_else(overflow)?; + let combined_term = combined_growth + .checked_add(combined_floor) + .ok_or_else(overflow)?; + + let selector_bound = attached_term + .checked_add(clone_term) + .and_then(|sum| sum.checked_add(combined_term)) + .ok_or_else(overflow)?; + selector_bound + .checked_mul(size_of::()) + .ok_or_else(overflow) +} + +/// Upper bound, in [`RowSelector`]s, on how many extra selectors the parquet reader's +/// page-index pruning can add on top of the deletion vector's own selection when normalizing +/// one file, from that file's already-fetched [`ParquetMetaData`]. +/// +/// `intersect_row_selections` (parquet's `selection.rs`), which combines a page-pruning +/// selection with the deletion vector's selection, is a `from_fn` generator whose +/// `size_hint()` is `(0, None)`: for inputs of length `a` and `b`, its output can have up to +/// `a + b` selectors -- longer than either input. Bounding the page-pruning side of that sum +/// requires knowing how many selectors a page-index-derived selection could produce: at most +/// two per data page (one skip, one select, in the worst case of alternating page-level +/// pruning decisions), summed over every column of every row group. +/// +/// Returns `0` when `metadata` carries no offset index (`metadata.offset_index()` is `None`). +/// This is provably safe, not merely a convenient default: page-index pruning cannot produce a +/// page-level selection without the offset index to locate pages by, so there are no +/// page-pruning selectors to bound. The offset index is fetched with +/// `PageIndexPolicy::Optional` from the same `FileMetadataCache` entry the scan's reader later +/// reopens (see [`attach_access_plan`]'s footer-fetch comment), so this function observes +/// exactly what the reader will see. +/// +/// Uses checked arithmetic throughout for the same reason as [`admission_bound_bytes`]. +fn page_selection_bound_selectors(metadata: &ParquetMetaData) -> Result { + let Some(offset_index) = metadata.offset_index() else { + return Ok(0); + }; + let overflow = || { + GeneralError( + "Deletion vector page-selection bound overflowed while summing offset-index page \ + locations" + .to_string(), + ) + }; + let mut total_page_locations = 0usize; + for row_group in offset_index { + for column in row_group { + total_page_locations = total_page_locations + .checked_add(column.page_locations().len()) + .ok_or_else(overflow)?; + } + } + total_page_locations.checked_mul(2).ok_or_else(overflow) +} + +/// Execution-memory-pool admission bound, in bytes, for one file's deletion-vector access +/// plan -- reserved *before* calling `build_access_plan` (see [`attach_access_plan`]'s +/// pre-reserve call site) to cover the larger of two peaks live at different points in the +/// plan's lifetime. In practice the reader-normalization peak below dominates the construction +/// peak unconditionally for any non-trivial input (`reader_peak_bytes(S, G) = (5S + 10G) * +/// size_of::()` always exceeds `CONSTRUCTION_PEAK_FACTOR * S * +/// size_of::() = 3S * size_of::()` once `S >= 1`, since the `5S` term +/// alone already exceeds `3S`); the construction term is retained as a documented floor rather +/// than dropped, since it is cheap to compute and keeps this bound correct even if the reader's +/// growth factors ever shrink below construction's. +/// +/// - **Construction peak** (`CONSTRUCTION_PEAK_FACTOR * S`, see that constant's doc comment): +/// live while `build_access_plan` builds the plan's `RowSelection`s. Construction's +/// transient allocations fully unwind before `build_access_plan` returns, so this peak never +/// overlaps the reader-normalization peak below. +/// - **Reader-normalization peak** (`reader_peak_bytes(S + page_bound_selectors, +/// num_row_groups)`, see that function): live later, once DataFusion's reader normalizes the +/// attached plan. `S = 2 * cardinality + num_row_groups` is the same conservative bound on +/// the plan's final retained selector count used for the construction peak -- it provably +/// bounds `R = total_selectors(&plan)` (`R <= S`, from `build_access_plan`'s +/// one-non-coalescing-selector-per-deleted-row worst case plus one boundary selector per row +/// group), so `S + page_bound_selectors` bounds `R` after page-index inflation the same way +/// `S` bounds `R` before it. +/// +/// These two peaks never overlap in time, so `max` -- not `sum` -- is the correct combinator: +/// reserving their sum would over-reserve for no safety benefit. +/// +/// Deliberately not clamped by the file's total row count here, unlike the reader-peak target +/// `attach_access_plan` resizes down to after construction (see that call site): `S`'s +/// `+ num_row_groups` boundary term is a worst-case padding margin that can legitimately exceed +/// the total row count for a small, heavily-deleted file, and admission sizing has no actual +/// retained-selector count yet to clamp against -- only after construction, once `R` is known, +/// is clamping to the total row count both meaningful and strictly tighter. Leaving this bound +/// unclamped only ever makes admission more conservative, never less safe. +/// +/// Uses checked arithmetic throughout: a cardinality, row-group count, or page bound large +/// enough to overflow `usize` while computing this bound indicates a corrupted or malicious +/// descriptor, reported as a clean error rather than panicking. +fn admission_bound_bytes( + cardinality: i64, + num_row_groups: usize, + page_bound_selectors: usize, +) -> Result { + let overflow = || { + GeneralError(format!( + "Deletion vector admission bound overflowed for cardinality {cardinality}, \ + {num_row_groups} row groups, and page bound {page_bound_selectors} selectors" + )) + }; + let cardinality_usize = usize::try_from(cardinality).map_err(|_| overflow())?; + // S: the conservative bound on the plan's final *retained* selector count (what + // `total_selectors(&plan)` cannot exceed) -- unchanged from the pre-existing + // construction-only bound this function replaces. + let s = cardinality_usize + .checked_mul(2) + .and_then(|doubled| doubled.checked_add(num_row_groups)) + .ok_or_else(overflow)?; + + let construction_bytes = s + .checked_mul(size_of::()) + .and_then(|bytes| bytes.checked_mul(CONSTRUCTION_PEAK_FACTOR)) + .ok_or_else(overflow)?; + + let s_plus_page = s.checked_add(page_bound_selectors).ok_or_else(overflow)?; + let reader_bytes = reader_peak_bytes(s_plus_page, num_row_groups)?; + + Ok(construction_bytes.max(reader_bytes)) +} + +/// Upper bound on concurrent DV-blob and footer fetches per partition. Both +/// are small ranged reads, so a modest fan-out hides object-store latency +/// without flooding the store client. +const DV_FETCH_CONCURRENCY: usize = 8; + +/// Called via `block_on` at plan-creation time on the executor task: DV blobs +/// are small ranged reads and footers are needed to learn row-group +/// boundaries. Files are fetched concurrently (bounded by +/// [`DV_FETCH_CONCURRENCY`]) with input order preserved. Footer fetches go +/// through the scan's shared FileMetadataCache, so the scan's subsequent open +/// of the same file is served from cache. That reuse relies on each input +/// [`PartitionedFile`] being returned as-is (only `with_extension` applied), +/// never rebuilt: the cache entry is keyed by this exact `object_meta` and the +/// scan later looks it up through the same struct. +/// +/// Deliberately takes no object-store options map and imports no +/// store-construction helper: every [`DvScanFile`] arrives with its stores +/// already resolved by the caller (see its doc comment), so this async path +/// structurally cannot build an object store -- only `runtime_env` is still +/// threaded through, for the shared `FileMetadataCache` and the execution +/// `MemoryPool` each expanded access plan's row selectors are reserved +/// against -- see [`DvAccessPlanReservation`]. +/// +/// Registers one memory consumer for the whole call and hands every DV'd file +/// its own empty reservation from that one registration. The per-file grow, +/// resize and release stay as they are, but the pool sees one consumer for the +/// partition rather than one per file. That matters for `CometFairMemoryPool`, +/// which divides the pool by the number of registered consumers: the +/// reservations live in the returned files until the task ends, so one +/// consumer per file would lower the fair limit of every other native operator +/// in the task, including a hash join build that cannot spill. The registration +/// itself is dropped when the last file holding a reservation from it is +/// dropped. When no file carries a DV, it is dropped when this call returns. +pub async fn attach_access_plans( + runtime_env: Arc, + files: Vec, +) -> Result, ExecutionError> { + let partition_reservation = + MemoryConsumer::new("DeltaDeletionVectorAccessPlan").register(&runtime_env.memory_pool); + futures::stream::iter(files) + .map(|scan_file| { + attach_access_plan(Arc::clone(&runtime_env), &partition_reservation, scan_file) + }) + .buffered(DV_FETCH_CONCURRENCY) + .try_collect() + .await +} + +/// Resolve one file's deletion vector into an attached [`ParquetAccessPlan`]; +/// files without a DV pass through untouched. `partition_reservation` is the +/// call's one registered consumer, and a DV'd file takes an empty reservation +/// from it rather than registering a consumer of its own. +async fn attach_access_plan( + runtime_env: Arc, + partition_reservation: &MemoryReservation, + scan_file: DvScanFile, +) -> Result { + let DvScanFile { + file, + file_path, + dv, + data_store, + dv_store, + } = scan_file; + let dv = match dv { + Some(dv) => dv, + None => return Ok(file), + }; + // Delta's canonical `DeletionVectorDescriptor.EMPTY`: inline storage, empty + // payload, size 0, cardinality 0. Spark's reader returns all rows for it; + // decoding would fail (the empty payload is too short for a magic + // number), so pass the file through unchanged before attempting to read it. + if dv.cardinality == 0 && dv.size_in_bytes == 0 { + return Ok(file); + } + if dv.size_in_bytes < 0 { + return Err(GeneralError(format!( + "Deletion vector for {file_path} has negative size {}", + dv.size_in_bytes + ))); + } + if dv.cardinality < 0 { + return Err(GeneralError(format!( + "Deletion vector for {file_path} has negative cardinality {}", + dv.cardinality + ))); + } + + let data: Vec = if let Some(inline) = dv.inline_data { + check_inline_payload_size(&file_path, &inline, dv.size_in_bytes)?; + inline + } else if let Some(dv_path) = &dv.absolute_path { + let offset = dv + .offset + .ok_or_else(|| GeneralError("On-disk deletion vector missing offset".into()))?; + if offset < 0 { + return Err(GeneralError(format!( + "Deletion vector for {file_path} has negative offset {offset}" + ))); + } + let offset = offset as u64; + // [i32 BE size][data: size_in_bytes][i32 BE crc] + let framed_len = 4 + dv.size_in_bytes as u64 + 4; + let (store, dv_store_path) = dv_store.ok_or_else(|| { + GeneralError(format!( + "Deletion vector for {file_path} has an absolute path but no pre-resolved object store" + )) + })?; + let blob = store + .get_range(&dv_store_path, offset..offset + framed_len) + .await + .map_err(|e| GeneralError(format!("Failed to read deletion vector {dv_path}: {e}")))?; + unframe_dv_blob(&blob, dv.size_in_bytes as usize)?.to_vec() + } else { + return Err(GeneralError( + "Deletion vector descriptor has neither inline data nor a path".into(), + )); + }; + let deleted = deserialize_dv_bitmap(&data) + .map_err(|e| GeneralError(format!("Invalid deletion vector for {file_path}: {e}")))?; + validate_cardinality(&file_path, dv.cardinality, &deleted)?; + + // Row-group boundaries come from the data file's footer, fetched through the scan's + // shared FileMetadataCache with the page index loaded eagerly and the scan's metadata + // size hint (mirroring EagerPageIndexReaderFactory): the one fetch here also serves the + // subsequent data-file open, so DV files pay no extra footer round-trip. Keyed by + // `file.object_meta`, the exact ObjectMeta the scan's reader factory will look up. + let metadata_cache = runtime_env.cache_manager.get_file_metadata_cache(); + let metadata = DFParquetMetadata::new(data_store.as_ref(), &file.object_meta) + .with_file_metadata_cache(Some(metadata_cache)) + .with_page_index_policy(Some(PageIndexPolicy::Optional)) + .with_metadata_size_hint(Some(crate::parquet::parquet_exec::METADATA_SIZE_HINT)) + .fetch_metadata() + .await + .map_err(|e| GeneralError(format!("Failed to read parquet footer of {file_path}: {e}")))?; + let row_counts: Vec = metadata + .row_groups() + .iter() + .map(|rg| rg.num_rows()) + .collect(); + + // Pre-reserve the admission bound *before* calling build_access_plan: this bound covers + // both construction's own transient peak AND the larger peak DataFusion's reader hits + // later while normalizing the attached plan (`create_initial_plan`'s deep clone plus + // `into_overall_row_selection`'s combined RowSelection) -- see admission_bound_bytes and + // reader_peak_bytes. Reserving first means a rejection happens before any large `Vec` is + // allocated, not after -- see reader_peak_bytes's doc comment for the measured worst + // cases. The error message names this as a construction-phase rejection (contains + // "construct"), textually distinct from the steady-state message below, so callers/logs + // can tell which phase failed. + let page_bound_selectors = page_selection_bound_selectors(&metadata)?; + let admission_bytes = + admission_bound_bytes(dv.cardinality, row_counts.len(), page_bound_selectors)?; + let reservation = partition_reservation.new_empty(); + reservation.try_grow(admission_bytes).map_err(|e| { + GeneralError(format!( + "Deletion vector access plan for {file_path} needs up to {admission_bytes} \ + bytes to construct, exceeding the execution memory pool: {e}" + )) + })?; + + let plan = build_access_plan(&metadata, &deleted, file.range.as_ref()) + .map_err(|e| GeneralError(format!("Invalid deletion vector for {file_path}: {e}")))?; + + // Shrink the reservation to the reader-lifecycle steady state now that construction's + // transient peak has passed: the peak DataFusion's reader hits later while normalizing + // this file's attached plan (see reader_peak_bytes), not merely the plan's own retained + // selector bytes. `Rp_bound` bounds the selector count the reader will see after + // page-index pruning inflates the deletion vector's own selection: this plan's actual + // retained selector count (`R = total_selectors(&plan)`) plus `page_bound_selectors`, + // clamped to the file's total row count -- a RowSelection can never carry more than one + // selector per row, so `total_rows` independently bounds the reader's true selector count + // regardless of how loose `R + page_bound_selectors` is. + // + // NEVER-GROWS PROOF (this call always shrinks -- never fails): `R <= S` (established by + // `build_access_plan`'s worst case, the same invariant `admission_bound_bytes` relies on + // for its own `S`), so `Rp_bound = min(R + page_bound_selectors, total_rows) <= + // R + page_bound_selectors <= S + page_bound_selectors` -- the exact quantity + // `admission_bound_bytes` fed into `reader_peak_bytes` when computing the reservation + // already made above. `reader_peak_bytes` is monotone non-decreasing in its first + // argument (all three terms of its sum scale with `selectors`, `num_row_groups`, or + // both), so + // `reader_peak_bytes(Rp_bound, num_row_groups) <= + // reader_peak_bytes(S + page_bound_selectors, num_row_groups) <= admission_bytes`. + // `try_resize` is still used (rather than the infallible `resize`) so a violation of that + // invariant surfaces as a clean error instead of an internal panic. + let selector_count = total_selectors(&plan); + let total_rows: usize = row_counts + .iter() + .try_fold(0usize, |sum, &n| { + usize::try_from(n).ok().and_then(|n| sum.checked_add(n)) + }) + .ok_or_else(|| { + GeneralError(format!( + "Deletion vector total row count negative or overflowed usize for {file_path}" + )) + })?; + let reader_selector_bound = selector_count + .checked_add(page_bound_selectors) + .ok_or_else(|| { + GeneralError(format!( + "Deletion vector reader-peak bound overflowed for {file_path} while adding the \ + page-index inflation term" + )) + })? + .min(total_rows); + let retained_bytes_bound = reader_peak_bytes(reader_selector_bound, row_counts.len())?; + reservation.try_resize(retained_bytes_bound).map_err(|e| { + GeneralError(format!( + "Deletion vector access plan for {file_path} retains {selector_count} row \ + selectors, needing up to {retained_bytes_bound} bytes at the reader's \ + normalization peak, exceeding the execution memory pool: {e}" + )) + })?; + + // Keyed by concrete type: the parquet opener looks up + // `extensions.get::()`, so the plan must be stored + // as ParquetAccessPlan itself, NOT wrapped in an Arc (which would key + // it as Arc and silently skip DV application). The + // reservation occupies its own slot (`DvAccessPlanReservation`, keyed + // separately by its own concrete type) alongside it -- `extensions` is + // a multi-slot, type-keyed map (`datafusion_common::extensions`), not a + // single-slot table, so the two coexist without conflict and are + // dropped together. + Ok(file + .with_extension(plan) + .with_extension(DvAccessPlanReservation(reservation))) +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::datatypes::Schema; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool}; + use datafusion::execution::runtime_env::RuntimeEnvBuilder; + use parquet::arrow::ArrowWriter; + use parquet::file::metadata::ParquetMetaDataReader; + use parquet::file::properties::WriterProperties; + + fn metadata_with_page_offsets( + row_counts: &[i64], + offsets: &[(i64, Option)], + ) -> ParquetMetaData { + use parquet::basic::Type; + use parquet::file::metadata::{ColumnChunkMetaData, FileMetaData, RowGroupMetaData}; + use parquet::schema::types::{SchemaDescriptor, Type as SchemaType}; + + assert_eq!(row_counts.len(), offsets.len()); + let field = Arc::new( + SchemaType::primitive_type_builder("id", Type::INT64) + .build() + .unwrap(), + ); + let schema = Arc::new( + SchemaType::group_type_builder("schema") + .with_fields(vec![field]) + .build() + .unwrap(), + ); + let descriptor = Arc::new(SchemaDescriptor::new(schema)); + let row_groups = row_counts + .iter() + .zip(offsets) + .map(|(&count, &(data_offset, dictionary_offset))| { + RowGroupMetaData::builder(Arc::clone(&descriptor)) + .set_num_rows(count) + .set_column_metadata(vec![ColumnChunkMetaData::builder(descriptor.column(0)) + .set_data_page_offset(data_offset) + .set_dictionary_page_offset(dictionary_offset) + .build() + .unwrap()]) + .build() + .unwrap() + }) + .collect(); + // Tests validate row-group counts, including corruption; do not sum them here. + let file_metadata = FileMetaData::new(1, 0, None, None, descriptor, None); + ParquetMetaData::new(file_metadata, row_groups) + } + + fn full_file_access_plan( + row_counts: &[i64], + deleted: &RoaringTreemap, + ) -> Result { + build_access_plan( + &metadata_with_page_offsets(row_counts, &vec![(0, None); row_counts.len()]), + deleted, + None, + ) + } + + /// Mirror the pre-resolution `plan_delta_spark_scan` does before entering + /// `attach_access_plans`: resolve `url`'s object store and within-store + /// path via the same helper the production code path uses, outside any + /// async runtime, exactly as `DvScanFile` requires. + fn resolve_store(runtime_env: &Arc, url: &str) -> (Arc, Path) { + use crate::parquet::parquet_support::prepare_object_store_with_configs; + let (store_url, path, _) = prepare_object_store_with_configs( + Arc::clone(runtime_env), + url.to_string(), + &std::collections::HashMap::new(), + ) + .unwrap(); + let store = runtime_env.object_store(&store_url).unwrap(); + (store, path) + } + + /// Build a one-column (`id: Int64`), `num_rows`-row batch (values `0..num_rows`), shared by + /// every parquet-writing helper below. + fn sequential_int64_batch(num_rows: i64) -> (Arc, RecordBatch) { + use datafusion::arrow::array::Int64Array; + use datafusion::arrow::datatypes::{DataType, Field}; + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from_iter_values(0..num_rows))], + ) + .unwrap(); + (schema, batch) + } + + /// Write a one-column parquet file with rows 0..num_rows using explicit `props`; returns + /// its size. + fn write_parquet_with_properties( + path: &std::path::Path, + num_rows: i64, + props: WriterProperties, + ) -> i64 { + let (schema, batch) = sequential_int64_batch(num_rows); + let out = std::fs::File::create(path).unwrap(); + let mut writer = ArrowWriter::try_new(out, schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + std::fs::metadata(path).unwrap().len() as i64 + } + + /// Write a one-column parquet file with rows 0..num_rows; returns its size. + fn write_parquet(path: &std::path::Path, num_rows: i64) -> i64 { + write_parquet_with_properties(path, num_rows, WriterProperties::default()) + } + + /// Write a two-row-group parquet file (`2 * rows_per_group` total rows, split evenly via + /// an explicit `max_row_group_size`); returns its size. Used by tests exercising + /// `into_overall_row_selection`'s per-`Scan`-group boundary-selector term. + fn write_two_row_groups(path: &std::path::Path, rows_per_group: i64) -> i64 { + write_parquet_with_properties( + path, + rows_per_group * 2, + WriterProperties::builder() + .set_max_row_group_row_count(Some(rows_per_group as usize)) + .build(), + ) + } + + /// Read `path`'s full [`ParquetMetaData`], including the page index, exactly as this + /// module's own footer fetch does (`PageIndexPolicy::Optional`) -- synchronously, for test + /// setup that needs the real metadata before entering `attach_access_plans`' async path. + fn read_metadata_with_page_index(path: &std::path::Path) -> ParquetMetaData { + let file = std::fs::File::open(path).unwrap(); + ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Optional) + .parse_and_finish(&file) + .unwrap() + } + + /// End-to-end over local files: inline and on-disk DVs resolve to attached + /// access plans, non-DV files pass through untouched, and the output keeps + /// the input's file order (which concurrent fetching must preserve). + #[tokio::test] + async fn attach_access_plans_resolves_dvs_and_preserves_order() { + let tmp = tempfile::tempdir().unwrap(); + let dir = tmp.path(); + + let inline_deleted: RoaringTreemap = [0u64].into_iter().collect(); + let inline_data = portable_bytes(&inline_deleted); + + // On-disk DV file: 1 version byte, then two framed blobs back to back, so the second + // one exercises the offset..offset+framed_len slice past the first. + let ondisk_deleted: RoaringTreemap = [1u64].into_iter().collect(); + let ondisk_data = portable_bytes(&ondisk_deleted); + let second_deleted: RoaringTreemap = [3u64].into_iter().collect(); + let second_data = portable_bytes(&second_deleted); + let dv_file = dir.join("dv.bin"); + let mut dv_bytes = vec![1u8]; + dv_bytes.extend(frame(&ondisk_data)); + let second_offset = dv_bytes.len() as i32; + dv_bytes.extend(frame(&second_data)); + std::fs::write(&dv_file, &dv_bytes).unwrap(); + + let dv_for = |name: &str| match name { + "f0" => Some(DeltaSparkDvDescriptor { + storage_type: "i".to_string(), + absolute_path: None, + inline_data: Some(inline_data.clone()), + offset: None, + size_in_bytes: inline_data.len() as i32, + cardinality: 1, + }), + "f2" => Some(DeltaSparkDvDescriptor { + storage_type: "p".to_string(), + absolute_path: Some(format!("file://{}", dv_file.display())), + inline_data: None, + offset: Some(1), + size_in_bytes: ondisk_data.len() as i32, + cardinality: 1, + }), + "f5" => Some(DeltaSparkDvDescriptor { + storage_type: "p".to_string(), + absolute_path: Some(format!("file://{}", dv_file.display())), + inline_data: None, + offset: Some(second_offset), + size_in_bytes: second_data.len() as i32, + cardinality: 1, + }), + // Delta's `DeletionVectorDescriptor.EMPTY`: inline storage, empty + // payload, size 0, cardinality 0. Must pass through unchanged + // without attempting to decode the (empty) payload. + "f4" => Some(DeltaSparkDvDescriptor { + storage_type: "i".to_string(), + absolute_path: None, + inline_data: Some(vec![]), + offset: None, + size_in_bytes: 0, + cardinality: 0, + }), + _ => None, + }; + + let runtime_env = Arc::new(RuntimeEnv::default()); + let names = ["f0", "f1", "f2", "f3", "f4", "f5"]; + let files: Vec = names + .iter() + .map(|name| { + let path = dir.join(format!("{name}.parquet")); + let size = write_parquet(&path, 10); + let file_path = format!("file://{}", path.display()); + let (data_store, _) = resolve_store(&runtime_env, &file_path); + let dv = dv_for(name); + let dv_store = dv + .as_ref() + .and_then(|d| d.absolute_path.as_deref()) + .map(|dv_path| resolve_store(&runtime_env, dv_path)); + DvScanFile { + file: PartitionedFile::new(path.display().to_string(), size as u64), + file_path, + dv, + data_store, + dv_store, + } + }) + .collect(); + + let out = attach_access_plans(Arc::clone(&runtime_env), files) + .await + .unwrap(); + + assert_eq!(out.len(), names.len()); + for (file, name) in out.iter().zip(names) { + assert!( + file.object_meta + .location + .as_ref() + .ends_with(&format!("{name}.parquet")), + "output order broken: expected {name}, got {}", + file.object_meta.location + ); + let plan = file.extensions.get::(); + match name { + "f0" | "f2" | "f5" => { + let plan = plan.unwrap_or_else(|| panic!("{name} should carry an access plan")); + let deleted_row = match name { + "f0" => 0, + "f2" => 1, + _ => 3, + }; + match &plan.inner()[0] { + RowGroupAccess::Selection(sel) => { + let selectors: Vec = sel.clone().into(); + let expected = if deleted_row == 0 { + vec![RowSelector::skip(1), RowSelector::select(9)] + } else { + vec![ + RowSelector::select(deleted_row), + RowSelector::skip(1), + RowSelector::select(9 - deleted_row), + ] + }; + assert_eq!(selectors, expected, "{name}"); + } + other => panic!("{name}: expected selection, got {other:?}"), + } + } + _ => assert!(plan.is_none(), "{name} should have no access plan"), + } + } + + // Footer reads must go through the shared FileMetadataCache so the scan's + // subsequent open of the same file is served from cache instead of paying a + // second footer round-trip. Files without a DV read no footer at all. + let cache = runtime_env.cache_manager.get_file_metadata_cache(); + for (file, name) in out.iter().zip(names) { + let cached = cache.get(&file.object_meta.location); + match name { + "f0" | "f2" | "f5" => assert!( + cached.is_some(), + "{name}: DV footer read should populate the shared metadata cache" + ), + _ => assert!( + cached.is_none(), + "{name}: no-DV file should not have fetched a footer" + ), + } + } + } + + /// Serialize a treemap in Delta's portable RoaringBitmapArray format + /// (magic + RoaringTreemap wire form). + fn portable_bytes(deleted: &RoaringTreemap) -> Vec { + let mut data = PORTABLE_MAGIC.to_le_bytes().to_vec(); + deleted.serialize_into(&mut data).unwrap(); + data + } + + /// Serialize values in Delta's "native" RoaringBitmapArray format. + fn native_bytes(values: &[u64]) -> Vec { + use std::collections::BTreeMap; + let mut by_key: BTreeMap = BTreeMap::new(); + for v in values { + by_key + .entry((v >> 32) as u32) + .or_default() + .insert(*v as u32); + } + let max_key = by_key.keys().max().copied().unwrap_or(0); + let mut data = NATIVE_MAGIC.to_le_bytes().to_vec(); + data.extend(((max_key + 1) as i32).to_le_bytes()); + for key in 0..=max_key { + let bitmap = by_key.remove(&key).unwrap_or_default(); + let mut bytes = Vec::new(); + bitmap.serialize_into(&mut bytes).unwrap(); + data.extend((bytes.len() as i32).to_le_bytes()); + data.extend(bytes); + } + data + } + + fn frame(data: &[u8]) -> Vec { + let mut blob = (data.len() as i32).to_be_bytes().to_vec(); + blob.extend_from_slice(data); + blob.extend((crc32fast::hash(data) as i32).to_be_bytes()); + blob + } + + #[test] + fn portable_roundtrip_through_framing() { + let deleted: RoaringTreemap = [1u64, 5, 6, 7, 1000, (3u64 << 32) + 42] + .into_iter() + .collect(); + let blob = frame(&portable_bytes(&deleted)); + let data = unframe_dv_blob(&blob, blob.len() - 8).unwrap(); + let decoded = deserialize_dv_bitmap(data).unwrap(); + assert_eq!(decoded, deleted); + } + + #[test] + fn native_format_decodes() { + let values = [0u64, 2, 3, 100, (1u64 << 32) + 7]; + let decoded = deserialize_dv_bitmap(&native_bytes(&values)).unwrap(); + let expected: RoaringTreemap = values.into_iter().collect(); + assert_eq!(decoded, expected); + } + + #[test] + fn framing_rejects_bad_size_and_crc() { + let deleted: RoaringTreemap = [1u64, 2].into_iter().collect(); + let blob = frame(&portable_bytes(&deleted)); + let err = unframe_dv_blob(&blob, 3).unwrap_err(); + assert!(format!("{err}").contains("size mismatch")); + + let mut corrupted = blob.clone(); + let mid = corrupted.len() / 2; + corrupted[mid] ^= 0xFF; + let err = unframe_dv_blob(&corrupted, blob.len() - 8).unwrap_err(); + let msg = format!("{err}"); + assert!( + msg.contains("checksum") || msg.contains("size mismatch"), + "unexpected: {msg}" + ); + } + + #[test] + fn cardinality_mismatch_is_rejected() { + let deleted: RoaringTreemap = [1u64].into_iter().collect(); + let bytes = portable_bytes(&deleted); + let decoded = deserialize_dv_bitmap(&bytes).unwrap(); + + let err = validate_cardinality("f.parquet", 2, &decoded).unwrap_err(); + let msg = format!("{err}"); + assert!(msg.contains("cardinality"), "unexpected: {msg}"); + + validate_cardinality("f.parquet", 1, &decoded).unwrap(); + } + + #[test] + fn access_plan_scan_skip_and_selection() { + // Three row groups of 10 rows: group 0 untouched, group 1 fully + // deleted, group 2 rows 21..24 deleted (local 1..4). + let deleted: RoaringTreemap = (10u64..20).chain(21u64..24).collect(); + let plan = full_file_access_plan(&[10, 10, 10], &deleted).unwrap(); + assert_eq!(&plan.inner()[0], &RowGroupAccess::Scan); + assert_eq!(&plan.inner()[1], &RowGroupAccess::Skip); + match &plan.inner()[2] { + RowGroupAccess::Selection(sel) => { + let selectors: Vec = sel.clone().into(); + assert_eq!( + selectors, + vec![ + RowSelector::select(1), + RowSelector::skip(3), + RowSelector::select(6) + ] + ); + } + other => panic!("expected selection, got {other:?}"), + } + } + + #[test] + fn access_plan_rejects_out_of_range_rows() { + let deleted: RoaringTreemap = [5u64, 25].into_iter().collect(); + let err = full_file_access_plan(&[10, 10], &deleted).unwrap_err(); + assert!(format!("{err}").contains("only has 20 rows")); + } + + #[test] + fn access_plan_rejects_negative_row_count_reported_by_a_corrupt_footer() { + // A corrupt footer can report a negative row count for a row group. Round-trip through + // the real parquet-crate RowGroupMetaData builder (`into_builder`, reusing a real row + // group's own column metadata rather than a bare negative literal) to prove the guard + // fires on the exact shape a corrupt footer would produce, not just an arbitrary i64. + let tmp = tempfile::tempdir().unwrap(); + let path = tmp.path().join("f.parquet"); + write_parquet(&path, 10); + let metadata = read_metadata_with_page_index(&path); + let corrupted = metadata + .row_group(0) + .clone() + .into_builder() + .set_num_rows(-5) + .build() + .unwrap(); + let row_counts = vec![corrupted.num_rows()]; + + let deleted = RoaringTreemap::new(); + let err = full_file_access_plan(&row_counts, &deleted).unwrap_err(); + let msg = err.to_string(); + assert!(msg.contains("-5"), "expected the negative value: {msg}"); + assert!( + msg.contains("row group 0"), + "expected the row group index: {msg}" + ); + } + + #[test] + fn access_plan_selects_complement_row_count() { + // Random-ish pattern in one 100-row group: every 7th row deleted. + let deleted: RoaringTreemap = (0u64..100).filter(|i| i % 7 == 0).collect(); + let plan = full_file_access_plan(&[100], &deleted).unwrap(); + match &plan.inner()[0] { + RowGroupAccess::Selection(sel) => { + let selected: usize = sel.iter().filter(|s| !s.skip).map(|s| s.row_count).sum(); + let skipped: usize = sel.iter().filter(|s| s.skip).map(|s| s.row_count).sum(); + assert_eq!(selected + skipped, 100); + assert_eq!(skipped, deleted.len() as usize); + } + other => panic!("expected selection, got {other:?}"), + } + } + + #[test] + fn split_access_plan_preserves_file_global_positions() { + let metadata = + metadata_with_page_offsets(&[3, 3, 3], &[(10, None), (20, None), (30, None)]); + let plan = build_access_plan( + &metadata, + &RoaringTreemap::from_iter([1, 4, 7]), + Some(&FileRange { start: 20, end: 21 }), + ) + .unwrap(); + assert_eq!(plan.inner()[0], RowGroupAccess::Skip); + assert_eq!(plan.inner()[2], RowGroupAccess::Skip); + assert_eq!( + plan.inner()[1], + RowGroupAccess::Selection(RowSelection::from(vec![ + RowSelector::select(1), + RowSelector::skip(1), + RowSelector::select(1), + ])) + ); + } + + #[test] + fn split_access_plan_uses_half_open_dictionary_page_boundaries() { + let metadata = + metadata_with_page_offsets(&[3, 3, 3], &[(20, Some(10)), (25, Some(20)), (30, None)]); + let plan = build_access_plan( + &metadata, + &RoaringTreemap::from_iter([1]), + Some(&FileRange { start: 20, end: 30 }), + ) + .unwrap(); + assert_eq!( + plan.inner(), + &[ + RowGroupAccess::Skip, + RowGroupAccess::Scan, + RowGroupAccess::Skip + ] + ); + } + + #[test] + fn split_access_plan_validates_excluded_rows_and_row_groups() { + let metadata = + metadata_with_page_offsets(&[3, 3, 3], &[(10, None), (20, None), (30, None)]); + for range in [ + FileRange { start: 20, end: 21 }, + FileRange { start: 40, end: 50 }, + ] { + let err = + build_access_plan(&metadata, &RoaringTreemap::from_iter([1, 9]), Some(&range)) + .unwrap_err(); + assert!(err.to_string().contains("only has 9 rows")); + } + // Metadata must be valid even when the split excludes the corrupt row group. + let metadata = metadata_with_page_offsets(&[3, -1], &[(10, None), (20, None)]); + assert!(build_access_plan( + &metadata, + &RoaringTreemap::new(), + Some(&FileRange { start: 10, end: 11 }), + ) + .is_err()); + let metadata = metadata_with_page_offsets( + &[i64::MAX, i64::MAX, i64::MAX], + &[(10, None), (20, None), (30, None)], + ); + assert!(build_access_plan( + &metadata, + &RoaringTreemap::new(), + Some(&FileRange { start: 10, end: 11 }), + ) + .is_err()); + } + + #[test] + fn split_access_plan_seeks_across_present_and_absent_roaring_containers() { + for (first_group_rows, second_group_rows, marked_row) in [ + ((1_i64 << 32) + 3, 3, (1_u64 << 32) + 4), + ((1_i64 << 32) + 100, 1_i64 << 32, (2_u64 << 32) + 50), + ] { + let metadata = metadata_with_page_offsets( + &[first_group_rows, second_group_rows], + &[(10, None), (20, None)], + ); + let plan = build_access_plan( + &metadata, + &RoaringTreemap::from_iter([1, 2, marked_row]), + Some(&FileRange { start: 20, end: 21 }), + ) + .unwrap(); + assert_eq!(plan.inner()[0], RowGroupAccess::Skip); + let RowGroupAccess::Selection(selection) = &plan.inner()[1] else { + panic!("expected a selection in the second row group"); + }; + assert_eq!(selection.row_count(), second_group_rows as usize - 1); + assert_eq!(selection.skipped_row_count(), 1); + } + + // A corrupt position beyond an absent high-32-bit container remains an error. + let metadata = metadata_with_page_offsets(&[(1_i64 << 32) + 100], &[(10, None)]); + assert!(build_access_plan( + &metadata, + &RoaringTreemap::from_iter([1, (2_u64 << 32) + 50]), + Some(&FileRange { start: 20, end: 21 }), + ) + .is_err()); + } + + #[test] + fn split_access_plan_has_no_selectors_for_excluded_groups() { + let metadata = metadata_with_page_offsets( + &[1_000, 1_000, 1_000], + &[(10, None), (20, None), (30, None)], + ); + let deleted = (0..1_000).chain(2_000..3_000).step_by(2).collect(); + let plan = build_access_plan(&metadata, &deleted, Some(&FileRange { start: 20, end: 21 })) + .unwrap(); + assert_eq!(total_selectors(&plan), 0); + assert_eq!(plan.inner()[1], RowGroupAccess::Scan); + assert!(total_selectors(&build_access_plan(&metadata, &deleted, None).unwrap()) > 0); + } + + /// The confirmed worst case: deleting every even row leaves + /// no adjacent skips or selects to merge, so `build_access_plan` emits + /// one non-coalescing `RowSelector` per row of the group. + fn alternating_deleted(num_rows: u64) -> RoaringTreemap { + (0..num_rows).step_by(2).collect() + } + + #[test] + fn total_selectors_counts_one_per_row_for_alternating_bitmap() { + let deleted = alternating_deleted(1024); + let plan = full_file_access_plan(&[1024], &deleted).unwrap(); + assert_eq!(total_selectors(&plan), 1024); + } + + #[test] + fn total_selectors_ignores_scan_and_skip_row_groups() { + // Group 0 untouched (Scan), group 1 fully deleted (Skip): neither + // carries a RowSelection, so both must contribute zero selectors. + let deleted: RoaringTreemap = (10u64..20).collect(); + let plan = full_file_access_plan(&[10, 10], &deleted).unwrap(); + assert_eq!(total_selectors(&plan), 0); + } + + /// Writes one file's on-disk parquet data for a full-file, alternating-bitmap deletion + /// vector, returning its path, byte size, and deleted-row bitmap so callers needing the + /// file's on-disk metadata (to size a memory pool exactly, or to replay the real reader + /// path) can inspect it before building a [`DvScanFile`] from it. + fn write_alternating_parquet( + dir: &std::path::Path, + num_rows: i64, + ) -> (std::path::PathBuf, i64, RoaringTreemap) { + let deleted = alternating_deleted(num_rows as u64); + let path = dir.join("alternating.parquet"); + let size = write_parquet(&path, num_rows); + (path, size, deleted) + } + + /// Builds a [`DvScanFile`] with an inline deletion vector for an already-written parquet + /// file at `path`. + fn dv_scan_file_for_alternating( + runtime_env: &Arc, + path: &std::path::Path, + size: i64, + deleted: &RoaringTreemap, + ) -> DvScanFile { + let inline_data = portable_bytes(deleted); + let file_path = format!("file://{}", path.display()); + let (data_store, _) = resolve_store(runtime_env, &file_path); + DvScanFile { + file: PartitionedFile::new(path.display().to_string(), size as u64), + file_path, + dv: Some(DeltaSparkDvDescriptor { + storage_type: "i".to_string(), + absolute_path: None, + inline_data: Some(inline_data.clone()), + offset: None, + size_in_bytes: inline_data.len() as i32, + cardinality: deleted.len() as i64, + }), + data_store, + dv_store: None, + } + } + + #[tokio::test] + async fn split_access_plan_reservation_tracks_only_retained_selectors() { + let tmp = tempfile::tempdir().unwrap(); + let path = tmp.path().join("split.parquet"); + let size = write_parquet_with_properties( + &path, + 3_000, + WriterProperties::builder() + .set_max_row_group_row_count(Some(1_000)) + .build(), + ); + let metadata = read_metadata_with_page_index(&path); + let column = metadata.row_group(1).column(0); + let offset = column + .dictionary_page_offset() + .unwrap_or_else(|| column.data_page_offset()); + let pool: Arc = Arc::new(GreedyMemoryPool::new(10_000_000)); + let runtime_env = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build_arc() + .unwrap(); + let deleted = alternating_deleted(3_000); + let full = dv_scan_file_for_alternating(&runtime_env, &path, size, &deleted); + let mut split = dv_scan_file_for_alternating(&runtime_env, &path, size, &deleted); + split.file.range = Some(FileRange { + start: offset, + end: offset + 1, + }); + let files = attach_access_plans(runtime_env, vec![full, split]) + .await + .unwrap(); + let full_plan = files[0].extensions.get::().unwrap(); + let split_plan = files[1].extensions.get::().unwrap(); + assert_eq!(total_selectors(full_plan), 3_000); + assert_eq!(total_selectors(split_plan), 1_000); + let full_reserved = files[0] + .extensions + .get::() + .unwrap() + .0 + .size(); + let split_reserved = files[1] + .extensions + .get::() + .unwrap() + .0 + .size(); + // Keep the existing reader-clone and page-index inflation allowance. Only the actual + // DV selector count changes, not the proof or the whole-file admission bound. + let page_bound = page_selection_bound_selectors(&metadata).unwrap(); + assert_eq!( + split_reserved, + reader_peak_bytes((1_000 + page_bound).min(3_000), 3).unwrap() + ); + assert!(split_reserved < full_reserved); + assert_eq!(pool.reserved(), full_reserved + split_reserved); + drop(files); + assert_eq!(pool.reserved(), 0); + } + + #[tokio::test] + async fn native_parquet_split_applies_global_deletions_with_and_without_pruning() { + use datafusion::arrow::array::Int64Array; + use datafusion::execution::object_store::ObjectStoreUrl; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; + use datafusion::physical_expr::PhysicalExpr; + use datafusion::physical_plan::ExecutionPlan; + use datafusion::prelude::{SessionConfig, SessionContext}; + use datafusion::scalar::ScalarValue; + + let tmp = tempfile::tempdir().unwrap(); + let path = tmp.path().join("split.parquet"); + let size = write_parquet_with_properties( + &path, + 30, + WriterProperties::builder() + .set_max_row_group_row_count(Some(10)) + .set_data_page_row_count_limit(2) + .set_write_batch_size(2) + .build(), + ); + let metadata = read_metadata_with_page_index(&path); + let column = metadata.row_group(1).column(0); + let offset = column + .dictionary_page_offset() + .unwrap_or_else(|| column.data_page_offset()); + let deleted = RoaringTreemap::from_iter([1, 4, 11, 14, 17, 21]); + for prune in [false, true] { + let runtime_env = RuntimeEnvBuilder::new().build_arc().unwrap(); + let mut file = dv_scan_file_for_alternating(&runtime_env, &path, size, &deleted); + file.file.range = Some(FileRange { + start: offset, + end: offset + 1, + }); + let files = attach_access_plans(Arc::clone(&runtime_env), vec![file]) + .await + .unwrap(); + let session = Arc::new(SessionContext::new_with_config_rt( + SessionConfig::new(), + runtime_env, + )); + let (schema, _) = sequential_int64_batch(0); + let filter: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("id", 0)), + Operator::GtEq, + Arc::new(Literal::new(ScalarValue::Int64(Some(15)))), + )); + let scan = crate::parquet::parquet_exec::init_datasource_exec( + Arc::clone(&schema), + Some(schema), + None, + ObjectStoreUrl::local_filesystem(), + crate::parquet::parquet_support::ObjectStoreBackend::Local, + vec![files], + None, + prune.then_some(vec![filter]), + None, + "UTC", + true, + false, + false, + false, + &session, + false, + false, + false, + false, + "", + "", + ) + .unwrap(); + let batches: Vec<_> = scan + .execute(0, session.task_ctx()) + .unwrap() + .try_collect() + .await + .unwrap(); + let values: Vec<_> = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect(); + let expected: Vec<_> = (10..20) + .filter(|id| !deleted.contains(*id as u64) && (!prune || *id >= 15)) + .collect(); + assert_eq!(values, expected, "pruning enabled: {prune}"); + if prune { + let Some(datafusion::physical_plan::metrics::MetricValue::PruningMetrics { + pruning_metrics, + .. + }) = scan + .metrics() + .unwrap() + .sum_by_name("page_index_rows_pruned") + else { + panic!("expected page-index pruning metrics"); + }; + assert!(pruning_metrics.pruned() > 0); + } + } + } + + /// Builds one file's [`DvScanFile`] carrying an inline, alternating-bitmap + /// deletion vector over `num_rows` -- enough retained selectors to make + /// the reservation's byte count non-trivial without needing an on-disk DV + /// file. Used by the memory-accounting tests below. + fn alternating_dv_scan_file( + runtime_env: &Arc, + dir: &std::path::Path, + num_rows: i64, + ) -> DvScanFile { + let (path, size, deleted) = write_alternating_parquet(dir, num_rows); + dv_scan_file_for_alternating(runtime_env, &path, size, &deleted) + } + + /// A pool too small for even one `RowSelector` must reject the file's + /// access plan with a clean, file-naming error instead of the caller + /// materializing the selectors unbounded and risking an executor OOM. + #[tokio::test] + async fn attach_access_plans_rejects_oversized_dv_against_tiny_pool() { + let tmp = tempfile::tempdir().unwrap(); + let pool: Arc = Arc::new(GreedyMemoryPool::new(1)); + let runtime_env = Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build() + .unwrap(), + ); + let scan_file = alternating_dv_scan_file(&runtime_env, tmp.path(), 1024); + + let err = attach_access_plans(Arc::clone(&runtime_env), vec![scan_file]) + .await + .unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("alternating.parquet"), + "error should name the file: {msg}" + ); + assert!( + msg.contains("Resources exhausted") || msg.contains("exceeding"), + "error should surface pool exhaustion: {msg}" + ); + assert!( + msg.to_lowercase().contains("construct"), + "a pool too small even for the construction-phase bound should fail with a \ + construction-phase message: {msg}" + ); + assert_eq!( + pool.reserved(), + 0, + "a rejected reservation must not leak bytes into the pool" + ); + } + + /// A pool with room for the plan succeeds, reserves exactly the reader-lifecycle peak + /// bound (`reader_peak_bytes`, never a hardcoded constant) once construction's transient + /// peak has passed, attaches the reservation alongside the access plan, and releases it + /// back to the pool when the returned files are dropped. + #[tokio::test] + async fn attach_access_plans_reserves_and_releases_selector_bytes() { + let tmp = tempfile::tempdir().unwrap(); + let pool: Arc = Arc::new(GreedyMemoryPool::new(1_000_000)); + let runtime_env = Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build() + .unwrap(), + ); + let scan_file = alternating_dv_scan_file(&runtime_env, tmp.path(), 1024); + // A full-file alternating bitmap retains exactly one selector per row (1024), which + // equals the file's total row count -- so the reader-peak clamp collapses to exactly + // this file's retained selector count regardless of its real page-index bound. + let expected_bytes = reader_peak_bytes(1024, 1).unwrap(); + + let out = attach_access_plans(Arc::clone(&runtime_env), vec![scan_file]) + .await + .unwrap(); + assert_eq!(out.len(), 1); + assert_eq!( + pool.reserved(), + expected_bytes, + "plan bytes should be reserved against the pool" + ); + + let reservation = out[0] + .extensions + .get::() + .expect("reservation extension should be attached alongside the access plan"); + assert_eq!(reservation.0.size(), expected_bytes); + + drop(out); + assert_eq!( + pool.reserved(), + 0, + "dropping the files should release the reservation back to the pool" + ); + } + + /// Multi-file variant of `attach_access_plans_reserves_and_releases_selector_bytes`: two + /// files with distinct alternating deletion vectors (different row counts, so distinct + /// selector byte counts) must have their reservations summed in the pool while the returned + /// files are alive, and released in full once every returned file is dropped. + #[tokio::test] + async fn attach_access_plans_reserves_and_releases_selector_bytes_for_multiple_files() { + let tmp_a = tempfile::tempdir().unwrap(); + let tmp_b = tempfile::tempdir().unwrap(); + let pool: Arc = Arc::new(GreedyMemoryPool::new(10_000_000)); + let runtime_env = Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build() + .unwrap(), + ); + let scan_file_a = alternating_dv_scan_file(&runtime_env, tmp_a.path(), 1024); + let scan_file_b = alternating_dv_scan_file(&runtime_env, tmp_b.path(), 512); + // Per-file sum: each full-file alternating bitmap's reader-peak bound is independent of + // the other file's row count (unlike a naive shared-factor formula would suggest). + let expected_bytes = + reader_peak_bytes(1024, 1).unwrap() + reader_peak_bytes(512, 1).unwrap(); + + let out = attach_access_plans(Arc::clone(&runtime_env), vec![scan_file_a, scan_file_b]) + .await + .unwrap(); + assert_eq!(out.len(), 2); + assert_eq!( + pool.reserved(), + expected_bytes, + "reserved bytes should be the SUM of both files' selector bytes while the files \ + are alive" + ); + + drop(out); + assert_eq!( + pool.reserved(), + 0, + "dropping the files should release every file's reservation back to the pool" + ); + } + + /// Mirrors `CometFairMemoryPool`'s admission check without a JVM: every registered + /// consumer counts toward the fair limit, `pool_size / registered`, and a grow is rejected + /// once the pool's total would exceed it. Also counts every `register` call so a test can + /// pin how many consumers one `attach_access_plans` call adds to the task. + #[derive(Debug)] + struct FairLimitPool { + pool_size: usize, + state: std::sync::Mutex, + } + + #[derive(Debug, Default)] + struct FairLimitState { + used: usize, + registered: usize, + register_calls: usize, + } + + impl FairLimitPool { + fn new(pool_size: usize) -> Self { + Self { + pool_size, + state: std::sync::Mutex::new(FairLimitState::default()), + } + } + + /// Consumers registered right now (register calls minus unregister calls). + fn registered(&self) -> usize { + self.state.lock().unwrap().registered + } + + /// Every `register` call ever made against this pool. + fn register_calls(&self) -> usize { + self.state.lock().unwrap().register_calls + } + } + + impl std::fmt::Display for FairLimitPool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let state = self.state.lock().unwrap(); + write!( + f, + "FairLimitPool(pool_size={}, used={}, registered={})", + self.pool_size, state.used, state.registered + ) + } + } + + impl MemoryPool for FairLimitPool { + fn name(&self) -> &str { + "FairLimitPool" + } + + fn register(&self, _: &MemoryConsumer) { + let mut state = self.state.lock().unwrap(); + state.registered += 1; + state.register_calls += 1; + } + + fn unregister(&self, _: &MemoryConsumer) { + let mut state = self.state.lock().unwrap(); + state.registered = state + .registered + .checked_sub(1) + .expect("unregister without a matching register"); + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.try_grow(reservation, additional).unwrap(); + } + + fn shrink(&self, _: &MemoryReservation, subtractive: usize) { + let mut state = self.state.lock().unwrap(); + state.used = state + .used + .checked_sub(subtractive) + .expect("shrink below the bytes tracked by the pool"); + } + + fn try_grow( + &self, + _: &MemoryReservation, + additional: usize, + ) -> datafusion::common::Result<()> { + if additional == 0 { + return Ok(()); + } + let mut state = self.state.lock().unwrap(); + let registered = state.registered; + let limit = self + .pool_size + .checked_div(registered) + .expect("try_grow with no registered consumer"); + let used = state.used; + if limit < used + additional { + return datafusion::common::resources_err!( + "Failed to acquire {additional} bytes where {used} bytes already reserved \ + and the fair limit is {limit} bytes, {registered} registered" + ); + } + state.used += additional; + Ok(()) + } + + fn reserved(&self) -> usize { + self.state.lock().unwrap().used + } + } + + /// One `attach_access_plans` call must add exactly one consumer to the task's pool no + /// matter how many DV'd files it attaches. `CometFairMemoryPool` divides the pool by the + /// number of registered consumers, and the reservations live in the returned files until + /// the task ends, so one consumer per file would lower every other native operator's fair + /// limit for the whole task even when the DV bytes themselves are tiny. Three DV'd files + /// must leave a later consumer (a hash join build, say) its half of the pool. Each file's + /// bytes must still return to the pool when that file alone drops, while the shared + /// registration stays until the last file is gone. + #[tokio::test] + async fn attach_access_plans_registers_one_consumer_per_call() { + let tmp_a = tempfile::tempdir().unwrap(); + let tmp_b = tempfile::tempdir().unwrap(); + let tmp_c = tempfile::tempdir().unwrap(); + let pool_size = 1_000_000usize; + let fair_pool = Arc::new(FairLimitPool::new(pool_size)); + let pool: Arc = Arc::clone(&fair_pool) as Arc; + let runtime_env = Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build() + .unwrap(), + ); + let scan_file_a = alternating_dv_scan_file(&runtime_env, tmp_a.path(), 64); + let scan_file_b = alternating_dv_scan_file(&runtime_env, tmp_b.path(), 128); + let scan_file_c = alternating_dv_scan_file(&runtime_env, tmp_c.path(), 256); + let bytes_a = reader_peak_bytes(64, 1).unwrap(); + let bytes_b = reader_peak_bytes(128, 1).unwrap(); + let bytes_c = reader_peak_bytes(256, 1).unwrap(); + + let mut out = attach_access_plans( + Arc::clone(&runtime_env), + vec![scan_file_a, scan_file_b, scan_file_c], + ) + .await + .unwrap(); + assert_eq!(out.len(), 3); + assert_eq!( + fair_pool.register_calls(), + 1, + "one attach_access_plans call should register exactly one consumer, not one per file" + ); + assert_eq!( + fair_pool.registered(), + 1, + "the one registration should stay while the returned files are alive" + ); + assert_eq!(pool.reserved(), bytes_a + bytes_b + bytes_c); + + // A second consumer in the same task sees a fair limit of half the pool. A third of + // the pool fits under that, and would not fit under a quarter or less. + let build_bytes = pool_size / 3; + assert!( + bytes_a + bytes_b + bytes_c + build_bytes <= pool_size / 2, + "test setup invariant: the DV bytes plus the build must fit under half the pool" + ); + let build = MemoryConsumer::new("HashJoinBuild").register(&pool); + build.try_grow(build_bytes).unwrap_or_else(|e| { + panic!("a second consumer should get its fair half of the pool: {e}") + }); + assert_eq!(fair_pool.registered(), 2); + + // Dropping one file returns only that file's bytes and keeps the shared registration. + let last = out.pop().unwrap(); + drop(last); + assert_eq!(pool.reserved(), bytes_a + bytes_b + build_bytes); + assert_eq!(fair_pool.registered(), 2); + + drop(out); + assert_eq!(pool.reserved(), build_bytes); + assert_eq!( + fair_pool.registered(), + 1, + "the shared registration should go once the last file drops" + ); + drop(build); + assert_eq!(pool.reserved(), 0); + assert_eq!(fair_pool.registered(), 0); + } + + /// A pool sized to fit only the larger of two files' selector bytes must reject the whole + /// batch -- regardless of which file's reservation attempt happens to run first under + /// `buffered`'s bounded concurrency -- and must not leave an earlier, transiently successful + /// file's reservation stranded in the pool once the batch's error propagates: `try_collect` + /// drops the whole in-flight `Vec` (including any already-resolved file's + /// attached `DvAccessPlanReservation`) as soon as any one file errors. + #[tokio::test] + async fn attach_access_plans_rejects_multi_file_batch_without_leaking_earlier_reservation() { + let tmp_a = tempfile::tempdir().unwrap(); + let tmp_b = tempfile::tempdir().unwrap(); + + // Write file A up front (rather than via `alternating_dv_scan_file`) so its on-disk + // metadata -- and thus its exact page-selection bound -- is available here, before the + // pool exists, to size `pool_capacity` using the exact same admission bound the + // production code computes. + let (path_a, size_a, deleted_a) = write_alternating_parquet(tmp_a.path(), 1024); + let metadata_a = read_metadata_with_page_index(&path_a); + let page_bound_a = page_selection_bound_selectors(&metadata_a).unwrap(); + + // Sized to exactly fit the larger file's (1024 rows, cardinality 512) admission bound + // alone -- derived, never hardcoded, so it tracks CONSTRUCTION_PEAK_FACTOR, + // reader_peak_bytes, and size_of::() across changes. Whichever of the two + // files reserves first (the FIRST reservation each file makes) fits alone, but the + // combined requirement (both files' admission bounds together) never does, so the + // batch fails no matter the scheduling order under `buffered`'s bounded concurrency. + let pool_capacity = admission_bound_bytes(512, 1, page_bound_a).unwrap(); + let pool: Arc = Arc::new(GreedyMemoryPool::new(pool_capacity)); + let runtime_env = Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build() + .unwrap(), + ); + let scan_file_a = dv_scan_file_for_alternating(&runtime_env, &path_a, size_a, &deleted_a); + let scan_file_b = alternating_dv_scan_file(&runtime_env, tmp_b.path(), 512); + + let err = attach_access_plans(Arc::clone(&runtime_env), vec![scan_file_a, scan_file_b]) + .await + .unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("Resources exhausted") || msg.contains("exceeding"), + "error should surface pool exhaustion: {msg}" + ); + assert_eq!( + pool.reserved(), + 0, + "a rejected multi-file batch must not leak bytes from any file's reservation, \ + including one that transiently succeeded before the batch as a whole failed" + ); + } + + /// A pool sized to fit only the STEADY-STATE reservation (`reader_peak_bytes` at this + /// file's actual retained selector count) but not the larger admission bound must still be + /// rejected: the pre-reserve step runs before `build_access_plan`, so undersizing only for + /// steady state is not enough to admit a file whose transient admission-phase peak the pool + /// cannot actually hold. The error must be textually distinguishable from a steady-state + /// rejection (contains "construct"). + #[tokio::test] + async fn construction_bound_rejects_before_building_the_plan() { + let num_rows = 1024i64; + let cardinality = 512i64; // alternating_deleted(1024).len() + let num_row_groups = 1usize; + + let tmp = tempfile::tempdir().unwrap(); + let (path, size, deleted) = write_alternating_parquet(tmp.path(), num_rows); + let metadata = read_metadata_with_page_index(&path); + let page_bound = page_selection_bound_selectors(&metadata).unwrap(); + + // A full-file alternating bitmap's actual retained selector count equals its total row + // count, so its reader-peak-clamped steady state is exactly reader_peak_bytes(num_rows, + // 1). This is strictly smaller than the admission bound below: S = 2 * cardinality + + // num_row_groups (1025) is strictly larger than num_rows == R (1024) for this file + // (S's one-selector row-group boundary padding), and the file's real page bound `P` + // (from its default-written offset index, `page_bound` above) further inflates the + // admission side via `S + P` -- so the true gap is + // `reader_peak_bytes(S + page_bound, 1) - reader_peak_bytes(num_rows, 1) == + // 5 * (S + page_bound - num_rows) * size_of::() == + // 5 * (1 + page_bound) * size_of::()`, not merely the 1-selector S/R + // difference alone. Deliberately near-tight, and NOT hardcoded to a specific byte + // count: `page_bound` is measured from the real file, not assumed to be zero. + let steady_state_bytes = reader_peak_bytes(num_rows as usize, num_row_groups).unwrap(); + let admission_bytes = + admission_bound_bytes(cardinality, num_row_groups, page_bound).unwrap(); + assert!( + steady_state_bytes < admission_bytes, + "test setup invariant: steady state ({steady_state_bytes}) must be smaller than the \ + admission bound ({admission_bytes}) for this rejection to be meaningful" + ); + + let pool: Arc = Arc::new(GreedyMemoryPool::new(steady_state_bytes)); + let runtime_env = Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build() + .unwrap(), + ); + let scan_file = dv_scan_file_for_alternating(&runtime_env, &path, size, &deleted); + + let err = attach_access_plans(Arc::clone(&runtime_env), vec![scan_file]) + .await + .unwrap_err(); + let msg = err.to_string(); + assert!( + msg.to_lowercase().contains("construct"), + "rejection at the pre-reserve step should carry a construction-phase message: {msg}" + ); + assert_eq!( + pool.reserved(), + 0, + "a rejected construction-phase reservation must not leak bytes into the pool" + ); + } + + /// Directly verifies the reader-peak invariant end to end: after `attach_access_plans` + /// completes, the attached reservation's steady-state size must equal `reader_peak_bytes` + /// evaluated at this file's actual retained selector count and row-group count -- + /// computed independently here via `build_access_plan`/`total_selectors`, never hardcoded + /// -- not the larger admission bound that was reserved up front. + #[tokio::test] + async fn steady_state_reservation_covers_reader_peak() { + let tmp = tempfile::tempdir().unwrap(); + let pool: Arc = Arc::new(GreedyMemoryPool::new(10_000_000)); + let runtime_env = Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build() + .unwrap(), + ); + let num_rows = 300i64; + let deleted = alternating_deleted(num_rows as u64); + let scan_file = alternating_dv_scan_file(&runtime_env, tmp.path(), num_rows); + + let out = attach_access_plans(Arc::clone(&runtime_env), vec![scan_file]) + .await + .unwrap(); + + let plan = full_file_access_plan(&[num_rows], &deleted).unwrap(); + let expected_bytes = reader_peak_bytes(total_selectors(&plan), 1).unwrap(); + + let reservation = out[0] + .extensions + .get::() + .expect("reservation extension should be attached alongside the access plan"); + assert_eq!(reservation.0.size(), expected_bytes); + assert_eq!(pool.reserved(), expected_bytes); + } + + #[test] + fn admission_bound_bytes_derives_from_cardinality_and_row_groups() { + let sel = size_of::(); + + // Zero cardinality, zero page bound: the reader term dominates in both cases below + // (reader_peak_bytes(S, G) = (5S + 10G) * sel always exceeds CONSTRUCTION_PEAK_FACTOR + // * S * sel = 3S * sel for S >= 1, since 5S alone already exceeds 3S). + assert_eq!( + admission_bound_bytes(0, 1, 0).unwrap(), + (CONSTRUCTION_PEAK_FACTOR * sel).max(reader_peak_bytes(1, 1).unwrap()) + ); + // Many row groups, still zero cardinality. + assert_eq!( + admission_bound_bytes(0, 1_000, 0).unwrap(), + (CONSTRUCTION_PEAK_FACTOR * 1_000 * sel).max(reader_peak_bytes(1_000, 1_000).unwrap()) + ); + // Typical case: cardinality dominates over a single row group, with a non-zero page + // bound feeding only the reader-normalization term. + let cardinality = 512usize; + let num_row_groups = 1usize; + let page_bound = 7usize; + let s = 2 * cardinality + num_row_groups; + assert_eq!( + admission_bound_bytes(cardinality as i64, num_row_groups, page_bound).unwrap(), + (CONSTRUCTION_PEAK_FACTOR * s * sel) + .max(reader_peak_bytes(s + page_bound, num_row_groups).unwrap()) + ); + + // Overflow anywhere in the derivation must produce a clean GeneralError, never a panic. + let err = admission_bound_bytes(0, usize::MAX, 0).unwrap_err(); + assert!(matches!(err, GeneralError(_)), "unexpected error: {err:?}"); + } + + /// Replays DataFusion 54.1's REAL reader-normalization path (not a reimplementation of + /// it): clones the attached plan exactly as `create_initial_plan` does, calls the actual, + /// public `ParquetAccessPlan::into_overall_row_selection` DataFusion will call from + /// `build_stream`, and recovers the resulting `RowSelection`'s TRUE backing `Vec` capacity + /// (not its length) -- the same quantity `reader_peak_bytes` bounds. Exercising the real + /// dependency rather than a model of it means this test keeps working (or fails loudly) + /// across future `datafusion`/`parquet` upgrades that change either crate's growth + /// strategy. + /// + /// This test (and `..._with_a_scan_row_group` below) covers the NO-page-index-pruning + /// path only: neither ever calls `scan_selection` on the clone, so `retained_selectors` + /// (from `total_selectors`, i.e. length, not capacity) is exact for BOTH the attached + /// original and the clone here -- see `reader_path_peak_fits_the_reservation_with_page_pruning` + /// for the case where the clone's own capacity can exceed its length. Also note the + /// assertion below is purely arithmetic: `attached_plan`, `cloned_plan`, and `combined` + /// are not necessarily all simultaneously resident in this process's memory at one program + /// point (Rust may reuse `cloned_plan`'s allocation once `into_overall_row_selection` + /// consumes it, before `combined` is bound) -- this test checks that the byte counts the + /// real dependency reports add up within the reservation, not that three buffers are + /// observed live at once via a profiler. + #[tokio::test] + async fn reader_path_peak_fits_the_reservation() { + let tmp = tempfile::tempdir().unwrap(); + let pool: Arc = Arc::new(GreedyMemoryPool::new(10_000_000)); + let runtime_env = Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build() + .unwrap(), + ); + let (path, size, deleted) = write_alternating_parquet(tmp.path(), 1024); + let scan_file = dv_scan_file_for_alternating(&runtime_env, &path, size, &deleted); + + let out = attach_access_plans(Arc::clone(&runtime_env), vec![scan_file]) + .await + .unwrap(); + let attached_plan = out[0] + .extensions + .get::() + .expect("attach_access_plans should have attached a plan") + .clone(); + let reservation = out[0] + .extensions + .get::() + .expect("reservation extension should be attached alongside the access plan"); + + let retained_selectors = total_selectors(&attached_plan); + + // Mirror create_initial_plan's deep clone: the original (still reachable via + // `out[0]`'s extensions) and the clone are live at once, exactly like the real reader. + let cloned_plan = attached_plan.clone(); + let metadata = read_metadata_with_page_index(&path); + let combined = cloned_plan + .into_overall_row_selection(metadata.row_groups()) + .unwrap() + .expect("a fully-alternating file should produce a combined RowSelection"); + // `From for Vec` moves the RowSelection's backing Vec, so + // this preserves its TRUE allocated capacity -- not merely its length. + let combined_selectors: Vec = combined.into(); + let combined_capacity = combined_selectors.capacity(); + + let peak_bytes = (retained_selectors + retained_selectors + combined_capacity) + * size_of::(); + assert!( + peak_bytes <= reservation.0.size(), + "the real DataFusion/parquet reader path's peak ({peak_bytes} bytes: \ + {retained_selectors} retained selectors x 2 live plan copies + \ + {combined_capacity} combined-selection Vec capacity) must fit the reservation \ + ({} bytes)", + reservation.0.size() + ); + } + + /// Same replay as `reader_path_peak_fits_the_reservation`, but with a two-row-group file + /// where only the first group has any deletions -- the second stays `RowGroupAccess::Scan` + /// (no `RowSelection`), exercising `into_overall_row_selection`'s one-`select`-per- + /// `Scan`-group term that a naive `k * total_selectors` bound would miss entirely. Like + /// that test, this one never calls `scan_selection` on the clone, so it exercises the + /// NO-page-index-pruning path only (clone length == clone capacity here); see the doc + /// comment there for why `retained_selectors` is exact in this test and why the assertion + /// below is arithmetic rather than a live-memory observation. + #[tokio::test] + async fn reader_path_peak_fits_the_reservation_with_a_scan_row_group() { + let tmp = tempfile::tempdir().unwrap(); + let pool: Arc = Arc::new(GreedyMemoryPool::new(10_000_000)); + let runtime_env = Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build() + .unwrap(), + ); + + // Two 500-row groups: only the first has any deletions, so the second stays a `Scan` + // row group in the resulting ParquetAccessPlan. + let rows_per_group = 500i64; + let deleted: RoaringTreemap = alternating_deleted(rows_per_group as u64); + let path = tmp.path().join("two_groups.parquet"); + let size = write_two_row_groups(&path, rows_per_group); + let scan_file = dv_scan_file_for_alternating(&runtime_env, &path, size, &deleted); + + let out = attach_access_plans(Arc::clone(&runtime_env), vec![scan_file]) + .await + .unwrap(); + let attached_plan = out[0] + .extensions + .get::() + .expect("attach_access_plans should have attached a plan") + .clone(); + assert_eq!( + &attached_plan.inner()[1], + &RowGroupAccess::Scan, + "the second, untouched row group must stay Scan" + ); + let reservation = out[0] + .extensions + .get::() + .expect("reservation extension should be attached alongside the access plan"); + + let retained_selectors = total_selectors(&attached_plan); + let cloned_plan = attached_plan.clone(); + let metadata = read_metadata_with_page_index(&path); + let combined = cloned_plan + .into_overall_row_selection(metadata.row_groups()) + .unwrap() + .expect("a plan with a Selection row group should produce a combined RowSelection"); + let combined_selectors: Vec = combined.into(); + let combined_capacity = combined_selectors.capacity(); + + let peak_bytes = (retained_selectors + retained_selectors + combined_capacity) + * size_of::(); + assert!( + peak_bytes <= reservation.0.size(), + "the real reader path's peak with a Scan row group present ({peak_bytes} bytes) \ + must fit the reservation ({} bytes)", + reservation.0.size() + ); + } + + /// Replays the page-index-pruning path that drives peak memory the highest: clones + /// the attached plan (mirroring `create_initial_plan`), then intersects the clone's + /// row-group `Selection` with a synthetic, all-selecting page `RowSelection` via + /// `ParquetAccessPlan::scan_selection` -- the EXACT call `access_plan.rs`'s row-group + /// intersection makes when `PagePruningAccessPlanFilter` fires + /// (`existing_selection.intersection(&page_derived)` -> `RowSelection::intersection` -> + /// `intersect_row_selections`, ANOTHER `from_fn` generator with `size_hint() == (0, + /// None)`). The synthetic selection selects every row of the row group (a no-op filter -- + /// it changes nothing about which rows are scanned), included ONLY to drive the clone + /// through the SAME capacity-inflating intersection path real page pruning takes, so the + /// recovered capacity reflects the real dependency's growth strategy, not a model of it. + /// `num_rows` is chosen just above a power of two (at test scale, `1,048,577` rows) so the + /// intersection's `next_power_of_two` capacity jump is real and visible, not accidentally + /// exact. + /// + /// Recovers BOTH the intersected clone's TRUE capacity and the subsequent combined + /// selection's TRUE capacity (each via `into_inner()` / pattern-matching by value and + /// `Into>`, never `.clone()` -- cloning a `RowSelection` resets capacity + /// to length, since `Vec::clone` allocates exactly `with_capacity(len)`), and asserts + /// `attached_len + clone_capacity + combined_capacity` fits the reservation. Unlike + /// `reader_path_peak_fits_the_reservation`, this test does NOT model the clone as exact -- + /// it is the one that would have caught the original under-count. + #[tokio::test] + async fn reader_path_peak_fits_the_reservation_with_page_pruning() { + let tmp = tempfile::tempdir().unwrap(); + let pool: Arc = Arc::new(GreedyMemoryPool::new(100_000_000)); + let runtime_env = Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build() + .unwrap(), + ); + let num_rows = 1025i64; // 2^10 + 1: next_power_of_two(1025) == 2048, a real jump. + let (path, size, deleted) = write_alternating_parquet(tmp.path(), num_rows); + let scan_file = dv_scan_file_for_alternating(&runtime_env, &path, size, &deleted); + + let out = attach_access_plans(Arc::clone(&runtime_env), vec![scan_file]) + .await + .unwrap(); + let attached_plan = out[0] + .extensions + .get::() + .expect("attach_access_plans should have attached a plan") + .clone(); + let reservation = out[0] + .extensions + .get::() + .expect("reservation extension should be attached alongside the access plan"); + let attached_len = total_selectors(&attached_plan); + + let all_select = RowSelection::from(vec![RowSelector::select(num_rows as usize)]); + + // Mirror create_initial_plan's clone, then simulate PagePruningAccessPlanFilter firing + // against it. + let mut clone_for_capacity = attached_plan.clone(); + clone_for_capacity.scan_selection(0, all_select.clone()); + // Recover the intersected clone's TRUE capacity: `into_inner()` moves the + // `Vec` out without cloning, and pattern-matching by value on the + // result moves the `RowSelection` out the same way -- neither step clones it. + let clone_selection = match clone_for_capacity.into_inner().into_iter().next().unwrap() { + RowGroupAccess::Selection(sel) => sel, + other => panic!( + "expected row group 0 to carry a Selection after scan_selection, got {other:?}" + ), + }; + let clone_selectors: Vec = clone_selection.into(); + let clone_capacity = clone_selectors.capacity(); + assert!( + clone_capacity > attached_len, + "test setup invariant: the intersection must actually inflate the clone's capacity \ + past its length ({attached_len}) for this test to exercise the fix -- got \ + {clone_capacity}" + ); + + // A second, independently-reconstructed intersected clone (identical content, so the + // SAME deterministic capacity) feeds into_overall_row_selection, mirroring how the + // real reader calls it on the plan AFTER page pruning has already mutated it in place. + let mut clone_for_combining = attached_plan.clone(); + clone_for_combining.scan_selection(0, all_select); + let metadata = read_metadata_with_page_index(&path); + let combined = clone_for_combining + .into_overall_row_selection(metadata.row_groups()) + .unwrap() + .expect("a plan with a Selection row group should produce a combined RowSelection"); + let combined_selectors: Vec = combined.into(); + let combined_capacity = combined_selectors.capacity(); + + let peak_bytes = + (attached_len + clone_capacity + combined_capacity) * size_of::(); + assert!( + peak_bytes <= reservation.0.size(), + "the real reader path's peak WITH page-index pruning firing against the clone \ + ({peak_bytes} bytes: {attached_len} attached selectors + {clone_capacity} \ + intersected-clone Vec capacity + {combined_capacity} combined-selection Vec \ + capacity) must fit the reservation ({} bytes)", + reservation.0.size() + ); + } + + /// Property check over a grid of `(cardinality, num_row_groups, page_bound)` combinations, + /// each checked at several `R <= S`: the reader-lifecycle steady-state bound can never + /// exceed the admission bound reserved up front -- the resize at the end of + /// `attach_access_plan` must never need to GROW the reservation, only shrink it. + #[test] + fn resize_never_grows() { + for cardinality in [0i64, 1, 5, 100, 1_000, 10_000] { + for num_row_groups in [1usize, 2, 5, 100] { + for page_bound in [0usize, 1, 3, 50] { + let s = 2 * cardinality as usize + num_row_groups; + let admission = + admission_bound_bytes(cardinality, num_row_groups, page_bound).unwrap(); + // Sample the real invariant `R <= S` at both extremes and the midpoint -- + // reader_peak_bytes is monotone in its first argument, so checking a few + // representative points is sufficient to catch a regression. + for &r in &[0usize, s / 2, s] { + let rp_bound = r + page_bound; + let reader_bytes = reader_peak_bytes(rp_bound, num_row_groups).unwrap(); + assert!( + reader_bytes <= admission, + "reader_peak_bytes({rp_bound}, {num_row_groups}) = {reader_bytes} \ + must not exceed admission_bound_bytes({cardinality}, \ + {num_row_groups}, {page_bound}) = {admission} for R={r} <= S={s}" + ); + } + } + } + } + } + + /// `page_selection_bound_selectors` must return exactly `0` when the file's metadata + /// carries no offset index (the `unwrap_or(0)` this module's doc comment claims is + /// provably safe, not merely a convenient default), and the shared + /// `PageIndexPolicy::Optional` fetch used throughout this module must actually populate the + /// offset index when the file has one -- otherwise every other test in this file exercising + /// `page_selection_bound_selectors` indirectly would be silently testing against `0` + /// instead of a real page-index bound. + #[test] + fn page_selection_bound_selectors_reflects_offset_index_presence() { + let tmp = tempfile::tempdir().unwrap(); + + // A file written with the offset index explicitly disabled: no page locations to bound. + let no_index_path = tmp.path().join("no_page_index.parquet"); + write_parquet_with_properties( + &no_index_path, + 1024, + WriterProperties::builder() + .set_offset_index_disabled(true) + .build(), + ); + let metadata_without_index = read_metadata_with_page_index(&no_index_path); + assert!( + metadata_without_index.offset_index().is_none(), + "test setup invariant: this file must have no offset index" + ); + assert_eq!( + page_selection_bound_selectors(&metadata_without_index).unwrap(), + 0 + ); + + // A file written with default properties: the offset index is written by default, and + // the PageIndexPolicy::Optional fetch this module uses must actually populate it. + let indexed_path = tmp.path().join("with_page_index.parquet"); + write_parquet(&indexed_path, 1024); + let metadata_with_index = read_metadata_with_page_index(&indexed_path); + assert!( + metadata_with_index.offset_index().is_some(), + "a default-written file should carry an offset index -- if this fails, the \ + Optional page-index fetch policy stopped populating it, and \ + page_selection_bound_selectors would be silently under-bounding" + ); + assert!( + page_selection_bound_selectors(&metadata_with_index).unwrap() > 0, + "a file with pages and an offset index should have a positive page-selection bound" + ); + } + + // ----------------------------------------------------------------------------------------- + // Malformed-input hardening matrix: every way a deletion-vector blob can be corrupted + // (truncation, CRC, magic, length lies, cardinality lies, and general bit-flip fuzzing) must + // yield a clean `Err`, NEVER a panic and never a silently wrong answer. + // ----------------------------------------------------------------------------------------- + + /// Runs `f` under `catch_unwind`, failing the test with `context` if it panics. Every + /// malformed-input case below routes through this so a panic surfaces as an attributable test + /// failure instead of aborting the whole test binary silently at whichever case triggered it. + fn assert_no_panic(context: &str, f: impl FnOnce() -> T + std::panic::UnwindSafe) -> T { + match std::panic::catch_unwind(f) { + Ok(result) => result, + Err(_) => panic!("panicked while decoding malformed input: {context}"), + } + } + + /// A valid on-disk-framed blob (`[i32 BE size][data][i32 BE crc]`) plus its unframed `data` + /// payload (the portable-format `[i32 LE magic][RoaringTreemap bytes]`, the same bytes an + /// inline DV descriptor would carry directly), shared by every malformed-input case below so + /// each corruption starts from one known-good baseline. + fn valid_dv_fixture() -> (Vec, Vec) { + let deleted: RoaringTreemap = [1u64, 5, 6, 7, 1000, (3u64 << 32) + 42] + .into_iter() + .collect(); + let data = portable_bytes(&deleted); + let blob = frame(&data); + (blob, data) + } + + /// An inline payload whose length disagrees with the descriptor's size must be rejected + /// before decoding; the JVM did the z85 decode, so this is native's only check point. + #[test] + fn inline_payload_length_must_match_descriptor_size() { + let (_blob, data) = valid_dv_fixture(); + let err = + check_inline_payload_size("part-0.parquet", &data, data.len() as i32 + 1).unwrap_err(); + let message = format!("{err}"); + assert!( + message.contains(&format!("{}", data.len())) + && message.contains(&format!("{}", data.len() + 1)), + "expected both lengths in: {message}" + ); + assert!(check_inline_payload_size("part-0.parquet", &data, data.len() as i32).is_ok()); + } + + /// Row counts that overflow when summed come from a corrupt footer; the sweep must fail + /// with the checked error rather than wrap and misfire the total-rows check. + #[test] + fn access_plan_rejects_row_counts_that_overflow_when_summed() { + let deleted: RoaringTreemap = [1u64].into_iter().collect(); + let err = full_file_access_plan(&[i64::MAX, i64::MAX, i64::MAX], &deleted).unwrap_err(); + assert!( + format!("{err}").contains("overflow"), + "expected an overflow error, got: {err}" + ); + } + + /// Deleting the last row of one group and the first row of the next lands one skip at + /// the tail of group k and one at the head of group k+1, with no selector crossing the + /// boundary. + #[test] + fn access_plan_handles_deleted_rows_on_a_row_group_boundary() { + let deleted: RoaringTreemap = [9u64, 10].into_iter().collect(); + let plan = full_file_access_plan(&[10, 10, 10], &deleted).unwrap(); + match &plan.inner()[0] { + RowGroupAccess::Selection(sel) => { + let selectors: Vec = sel.clone().into(); + assert_eq!( + selectors, + vec![RowSelector::select(9), RowSelector::skip(1)] + ); + } + other => panic!("expected selection in group 0, got {other:?}"), + } + match &plan.inner()[1] { + RowGroupAccess::Selection(sel) => { + let selectors: Vec = sel.clone().into(); + assert_eq!( + selectors, + vec![RowSelector::skip(1), RowSelector::select(9)] + ); + } + other => panic!("expected selection in group 1, got {other:?}"), + } + assert_eq!(&plan.inner()[2], &RowGroupAccess::Scan); + } + + /// (1) Truncating a valid on-disk-framed blob at EVERY byte length from 0 to `len - 1` must + /// be rejected cleanly by `unframe_dv_blob`, never panic -- covers every truncation point in + /// one deterministic sweep rather than a few hand-picked lengths. + #[test] + fn unframe_rejects_every_truncation_length() { + let (blob, data) = valid_dv_fixture(); + let expected_size = data.len(); + for len in 0..blob.len() { + let truncated = &blob[..len]; + let result = assert_no_panic(&format!("on-disk blob truncated to {len} bytes"), || { + unframe_dv_blob(truncated, expected_size) + }); + assert!( + result.is_err(), + "truncating the on-disk blob to {len}/{} bytes should be rejected", + blob.len() + ); + } + } + + /// (1, inline-DV path) `attach_access_plan` feeds an inline descriptor's `inline_data` + /// straight to `deserialize_dv_bitmap`, skipping `unframe_dv_blob` entirely -- it carries no + /// `[size][data][crc]` framing, just `[i32 LE magic]...`. Every truncation length of that + /// unframed payload must also be handled cleanly: either a clean `Err`, or -- if a truncated + /// prefix happens to still parse -- a well-formed treemap that `build_access_plan` can + /// consume without panicking. Never a panic in either step. + #[test] + fn deserialize_rejects_every_truncation_length_of_inline_payload() { + let (_blob, data) = valid_dv_fixture(); + for len in 0..data.len() { + let truncated = &data[..len]; + let context = format!("inline payload truncated to {len} bytes"); + let result = assert_no_panic(&context, || deserialize_dv_bitmap(truncated)); + if let Ok(treemap) = result { + let max_row = treemap + .max() + .and_then(|m| m.checked_add(1)) + .unwrap_or(u64::MAX); + assert_no_panic(&format!("{context}: build_access_plan on survivor"), || { + let _ = full_file_access_plan(&[max_row as i64], &treemap); + }); + } + } + } + + /// (2) Flipping each byte of the CRC field individually must be rejected as a checksum + /// mismatch. XORing with `0xFF` guarantees the flipped byte differs from its original value + /// at that position, so every flip actually corrupts the checksum -- it can never coincide + /// with the real value by construction. + #[test] + fn unframe_rejects_every_crc_byte_flip() { + let (blob, data) = valid_dv_fixture(); + let expected_size = data.len(); + let crc_start = blob.len() - 4; + for i in crc_start..blob.len() { + let mut corrupted = blob.clone(); + corrupted[i] ^= 0xFF; + let context = format!("CRC byte {i} flipped"); + let result = assert_no_panic(&context, || unframe_dv_blob(&corrupted, expected_size)); + let err = result.unwrap_err(); + assert!( + format!("{err}").contains("checksum"), + "{context} should be reported as a checksum mismatch: {err}" + ); + } + } + + /// (3) A magic number that matches neither known format must be rejected by name -- checked + /// against both obviously-wrong values and the bitwise complement of each real magic (which, + /// by construction, can never accidentally equal either real magic). + #[test] + fn deserialize_rejects_corrupted_magic() { + let (_blob, data) = valid_dv_fixture(); + let payload = &data[4..]; // magic-stripped body, reused under every corrupted magic + for bad_magic in [0i32, 1, -1, i32::MAX, !PORTABLE_MAGIC, !NATIVE_MAGIC] { + assert_ne!(bad_magic, PORTABLE_MAGIC); + assert_ne!(bad_magic, NATIVE_MAGIC); + let mut corrupted = bad_magic.to_le_bytes().to_vec(); + corrupted.extend_from_slice(payload); + let context = format!("magic corrupted to {bad_magic}"); + let result = assert_no_panic(&context, || deserialize_dv_bitmap(&corrupted)); + let err = result.unwrap_err(); + assert!( + format!("{err}").contains("magic"), + "{context}: unexpected error: {err}" + ); + } + } + + /// (4a) A declared size larger than the buffer actually holds must be rejected as truncated + /// -- not read out of bounds, not panic -- even when the descriptor's `expected_size` agrees + /// with the (lied-about) declared size, so it is the truncation check, not the size-mismatch + /// check, that has to catch it. + #[test] + fn unframe_rejects_size_field_larger_than_buffer() { + let (_blob, data) = valid_dv_fixture(); + let lie = data.len() + 1_000_000; // declares far more data than the buffer holds + let mut lied_blob = (lie as i32).to_be_bytes().to_vec(); + lied_blob.extend_from_slice(&data); + lied_blob.extend((crc32fast::hash(&data) as i32).to_be_bytes()); + let result = assert_no_panic("size field lies larger than the buffer", || { + unframe_dv_blob(&lied_blob, lie) + }); + let err = result.unwrap_err(); + assert!( + format!("{err}").contains("truncated"), + "unexpected error: {err}" + ); + } + + /// (4b) A declared size smaller than the data actually written shifts which bytes get hashed + /// as the CRC input, so it must surface as a checksum mismatch -- never a panic, never a + /// successful decode of a differently-sliced payload. + #[test] + fn unframe_rejects_size_field_smaller_than_actual_data() { + let (_blob, data) = valid_dv_fixture(); + let lie = data.len() - 4; // declares less data than was actually written + let mut lied_blob = (lie as i32).to_be_bytes().to_vec(); + lied_blob.extend_from_slice(&data); + lied_blob.extend((crc32fast::hash(&data) as i32).to_be_bytes()); + let result = assert_no_panic("size field lies smaller than actual data", || { + unframe_dv_blob(&lied_blob, lie) + }); + let err = result.unwrap_err(); + assert!( + format!("{err}").contains("checksum"), + "unexpected error: {err}" + ); + } + + /// (5) `validate_cardinality` must reject absurd cardinality claims in BOTH directions -- far + /// too high (a stale descriptor claiming millions of deletions for a handful of actual bits) + /// and far too low (zero or negative expected against many actual bits) -- and must never + /// panic, including when a negative `expected` (an `i64`) is cast to the `u64` comparison + /// `deleted.len()` uses. + #[test] + fn validate_cardinality_rejects_absurd_mismatches_in_both_directions() { + let deleted: RoaringTreemap = (0u64..1000).collect(); // 1000 actual deletions + + for (context, expected) in [ + ("claimed far too high", i64::MAX), + ("claimed far too low (zero)", 0i64), + ("claimed negative", -1i64), + ] { + let result = assert_no_panic(context, || { + validate_cardinality("f.parquet", expected, &deleted) + }); + let err = result.unwrap_err(); + assert!( + format!("{err}").contains("cardinality"), + "{context}: unexpected error: {err}" + ); + } + + // Reverse imbalance: an empty bitmap against a huge claimed cardinality. + let empty = RoaringTreemap::new(); + let result = assert_no_panic("empty bitmap vs huge claimed cardinality", || { + validate_cardinality("f.parquet", 1_000_000_000, &empty) + }); + let err = result.unwrap_err(); + assert!( + format!("{err}").contains("cardinality"), + "unexpected error: {err}" + ); + } + + /// (6) Single-bit-flip fuzz sweep over one valid on-disk-framed blob: for every bit position, + /// flip it and run the FULL decode pipeline (`unframe_dv_blob` then `deserialize_dv_bitmap`). + /// Every outcome must be either a clean `Err` or a successfully-decoded, well-formed treemap + /// that `build_access_plan` can consume without panicking -- NEVER a panic in either step. + /// Bounded to one pass over one blob's bits, so runtime stays well under 5s. + #[test] + fn single_bit_flip_sweep_never_panics() { + let (blob, data) = valid_dv_fixture(); + let expected_size = data.len(); + let mut checked = 0usize; + + for byte_idx in 0..blob.len() { + for bit in 0u8..8 { + let mut corrupted = blob.clone(); + corrupted[byte_idx] ^= 1 << bit; + let context = format!("byte {byte_idx} bit {bit} flipped"); + checked += 1; + + let unframed = assert_no_panic(&context, || { + unframe_dv_blob(&corrupted, expected_size).map(|d| d.to_vec()) + }); + let Ok(unframed_data) = unframed else { + continue; + }; + + let decoded = assert_no_panic(&context, || deserialize_dv_bitmap(&unframed_data)); + if let Ok(treemap) = decoded { + // A "VALID selection": consuming the decoded treemap downstream must not + // panic either, whatever its contents happen to be. `checked_add` avoids an + // overflow panic (rather than a clean Err) if corruption produced a max value + // of u64::MAX. + let max_row = treemap + .max() + .and_then(|m| m.checked_add(1)) + .unwrap_or(u64::MAX); + assert_no_panic(&format!("{context}: build_access_plan"), || { + let _ = full_file_access_plan(&[max_row as i64], &treemap); + }); + } + } + } + assert_eq!( + checked, + blob.len() * 8, + "every single-bit flip must have been exercised" + ); + } + + /// (6, inline-DV path) Same single-bit-flip sweep as above, but over the shorter, unframed + /// inline payload (`deserialize_dv_bitmap` only, no `unframe_dv_blob`) -- the exact bytes an + /// inline `DeltaSparkDvDescriptor.inline_data` carries. Bounded to one pass over one + /// (shorter) payload's bits. + #[test] + fn inline_payload_single_bit_flip_sweep_never_panics() { + let (_blob, data) = valid_dv_fixture(); + let mut checked = 0usize; + + for byte_idx in 0..data.len() { + for bit in 0u8..8 { + let mut corrupted = data.clone(); + corrupted[byte_idx] ^= 1 << bit; + let context = format!("inline byte {byte_idx} bit {bit} flipped"); + checked += 1; + + let decoded = assert_no_panic(&context, || deserialize_dv_bitmap(&corrupted)); + if let Ok(treemap) = decoded { + let max_row = treemap + .max() + .and_then(|m| m.checked_add(1)) + .unwrap_or(u64::MAX); + assert_no_panic(&format!("{context}: build_access_plan"), || { + let _ = full_file_access_plan(&[max_row as i64], &treemap); + }); + } + } + } + assert_eq!( + checked, + data.len() * 8, + "every single-bit flip of the inline payload must have been exercised" + ); + } +} diff --git a/native/core/src/execution/mod.rs b/native/core/src/execution/mod.rs index 9a2e54783b4..02228f5b563 100644 --- a/native/core/src/execution/mod.rs +++ b/native/core/src/execution/mod.rs @@ -18,6 +18,8 @@ //! PoC of vectorization execution through JNI to Rust. pub mod c_udf; pub mod columnar_to_row; +#[cfg(feature = "delta")] +pub mod delta_dv; pub mod expressions; pub mod jni_api; pub(crate) mod merge_as_partial; diff --git a/native/core/src/execution/operators/dynamic_filter/join/tests.rs b/native/core/src/execution/operators/dynamic_filter/join/tests.rs index fd2248a053f..8086ac647e8 100644 --- a/native/core/src/execution/operators/dynamic_filter/join/tests.rs +++ b/native/core/src/execution/operators/dynamic_filter/join/tests.rs @@ -670,6 +670,9 @@ fn parquet_probe( false, false, false, + false, + "CORRECTED", + "CORRECTED", ) .unwrap(); (file, scan) @@ -775,6 +778,9 @@ async fn reader_filter_crosses_null_check_conjunction_and_retains_residual() { false, false, false, + false, + "CORRECTED", + "CORRECTED", ) .unwrap(); let checks = [("key", 0), ("payload", 1), ("other", 2)].map(|(name, index)| { diff --git a/native/core/src/execution/operators/dynamic_filter/join/tests/schema_errors.rs b/native/core/src/execution/operators/dynamic_filter/join/tests/schema_errors.rs index 8fa236f23d8..f37b78ad2c6 100644 --- a/native/core/src/execution/operators/dynamic_filter/join/tests/schema_errors.rs +++ b/native/core/src/execution/operators/dynamic_filter/join/tests/schema_errors.rs @@ -86,6 +86,9 @@ fn scan( false, false, false, + false, + "CORRECTED", + "CORRECTED", ) .unwrap() } diff --git a/native/core/src/execution/operators/dynamic_filter/join/tests/schema_errors/partition_columns.rs b/native/core/src/execution/operators/dynamic_filter/join/tests/schema_errors/partition_columns.rs index 3a4d55ae06e..f6b00114f10 100644 --- a/native/core/src/execution/operators/dynamic_filter/join/tests/schema_errors/partition_columns.rs +++ b/native/core/src/execution/operators/dynamic_filter/join/tests/schema_errors/partition_columns.rs @@ -71,6 +71,9 @@ fn partitioned_scan( false, false, false, + false, + "CORRECTED", + "CORRECTED", ) .unwrap() } diff --git a/native/core/src/execution/operators/dynamic_filter/join/tests/timestamp_errors.rs b/native/core/src/execution/operators/dynamic_filter/join/tests/timestamp_errors.rs index 2da9af30cb5..aa982280517 100644 --- a/native/core/src/execution/operators/dynamic_filter/join/tests/timestamp_errors.rs +++ b/native/core/src/execution/operators/dynamic_filter/join/tests/timestamp_errors.rs @@ -105,6 +105,9 @@ async fn assert_timestamp_overflow_preserved(nested: bool) { false, false, false, + false, + "CORRECTED", + "CORRECTED", ) .unwrap(); let join = single_key_join_plans( diff --git a/native/core/src/execution/operators/dynamic_filter/topk/tests.rs b/native/core/src/execution/operators/dynamic_filter/topk/tests.rs index 46e72071386..58a9866bca0 100644 --- a/native/core/src/execution/operators/dynamic_filter/topk/tests.rs +++ b/native/core/src/execution/operators/dynamic_filter/topk/tests.rs @@ -196,6 +196,9 @@ fn parquet_scan( false, false, false, + false, + "CORRECTED", + "CORRECTED", ) .unwrap() } diff --git a/native/core/src/execution/operators/dynamic_filter/topk/tests/timestamp.rs b/native/core/src/execution/operators/dynamic_filter/topk/tests/timestamp.rs index cdefa621c43..ed5b7267e0a 100644 --- a/native/core/src/execution/operators/dynamic_filter/topk/tests/timestamp.rs +++ b/native/core/src/execution/operators/dynamic_filter/topk/tests/timestamp.rs @@ -95,6 +95,9 @@ fn timestamp_input( false, false, false, + false, + "CORRECTED", + "CORRECTED", ) .unwrap(); (file, scan) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 15a5c48acb3..1450220704e 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -31,6 +31,9 @@ mod write; // and calls into that crate. #[cfg(feature = "contrib-delta")] mod delta_scan; +// JVM-planned Delta sibling of the kernel handler above; see delta_spark_scan.rs. +#[cfg(feature = "delta")] +mod delta_spark_scan; #[cfg(feature = "contrib-lance")] mod lance_scan; @@ -99,13 +102,14 @@ use iceberg::expr::Bind; use crate::execution::operators::ExecutionError::GeneralError; use crate::execution::spark_plan::SparkPlan; use crate::parquet::objectstore::s3_blob_fs_support::normalize_object_store_url; -use crate::parquet::parquet_support::prepare_object_store_with_configs; +use crate::parquet::parquet_support::{prepare_object_store_with_configs, ObjectStoreBackend}; use datafusion::common::scalar::ScalarStructBuilder; use datafusion::common::{ tree_node::{Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeRewriter}, JoinType as DFJoinType, NullEquality, ScalarValue, }; use datafusion::datasource::listing::PartitionedFile; +use datafusion::datasource::object_store::ObjectStoreUrl; use datafusion::logical_expr::type_coercion::functions::fields_with_udf; use datafusion::logical_expr::{ AggregateUDF, ReturnFieldArgs, ScalarUDF, TypeSignature, WindowFrame, WindowFrameBound, @@ -510,6 +514,187 @@ impl PhysicalPlanner { self.partition } + /// Build the native parquet `DataSourceExec` shared by the parquet-backed scan arms + /// (NativeScan and, behind the `delta` feature, DeltaScan): schema conversion, data-filter + /// binding, object-store setup, file-group construction, and `init_datasource_exec`. + /// Arm-specific concerns (file-list decoding, deletion-vector handling) stay in the arms. + /// `rebase_from_file_metadata` opts the scan into per-file datetime calendar-rebase + /// resolution (see `datetime_rebase.rs`): the Delta arm passes true, while NativeScan + /// passes false to keep its documented no-rebase behavior (#5010). + /// `datetime_rebase_mode_in_read` / `int96_rebase_mode_in_read` carry the session's + /// effective read modes for files whose footer metadata does not decide the policy; + /// they are only consulted when `rebase_from_file_metadata` is true (empty means + /// EXCEPTION, the conservative refuse-ancient posture). + #[allow(clippy::too_many_arguments)] + fn build_parquet_scan_plan( + &self, + plan_id: u32, + common: &spark_operator::NativeScanCommon, + object_store_url: ObjectStoreUrl, + object_store_backend: ObjectStoreBackend, + files: Vec, + rebase_from_file_metadata: bool, + datetime_rebase_mode_in_read: &str, + int96_rebase_mode_in_read: &str, + ) -> Result, ExecutionError> { + let data_schema = convert_spark_types_to_arrow_schema(common.data_schema.as_slice()); + let required_schema: SchemaRef = + convert_spark_types_to_arrow_schema(common.required_schema.as_slice()); + let partition_schema: SchemaRef = + convert_spark_types_to_arrow_schema(common.partition_schema.as_slice()); + let projection_vector: Vec = common + .projection_vector + .iter() + .map(|offset| *offset as usize) + .collect(); + + // Check if this partition has any files (bucketed scan with bucket pruning may have + // empty partitions; a fully-pruned Delta partition likewise). + if files.is_empty() { + // Match the nonempty scan's projected table schema, including partition + // and constant metadata columns. Fused operators bind against this schema + // even when the bucket has no rows. + let partition_fields: Vec = partition_schema + .fields() + .iter() + .map(|f| Arc::new(Field::new(f.name(), f.data_type().clone(), f.is_nullable()))) + .collect(); + let table_schema = TableSchema::builder(data_schema) + .with_table_partition_cols(partition_fields) + .build(); + let output_schema = Arc::new(table_schema.table_schema().project(&projection_vector)?); + let empty_exec = Arc::new(EmptyExec::new(output_schema)); + return Ok(Arc::new(SparkPlan::new(plan_id, empty_exec, vec![]))); + } + + // data_filters may reference partition columns and constant metadata columns + // (e.g. `_metadata.file_size`), which the Parquet reader appends after + // required_schema's columns once partition_values are projected into the + // batch. Bind against the combined schema so `Bound` indices resolve + // correctly -- Scala's `exprToProto(filter, scan.output)` + // (CometNativeScan.scala) numbers columns against that same ordering. + let data_filters: Result>, ExecutionError> = + if common.data_filters.is_empty() { + Ok(vec![]) + } else { + let filter_schema: SchemaRef = Arc::new(Schema::new( + required_schema + .fields() + .iter() + .chain(partition_schema.fields().iter()) + .cloned() + .collect::>(), + )); + let float_operands = self.data_filter_float_operands(); + common + .data_filters + .iter() + .map(|expr| { + self.create_data_filter(expr, Arc::clone(&filter_schema), float_operands) + }) + .collect() + }; + + let default_values = self.parse_default_values(common, &required_schema)?; + + let file_groups: Vec> = vec![files]; + + let scan = init_datasource_exec( + required_schema, + Some(data_schema), + Some(partition_schema), + object_store_url, + object_store_backend, + file_groups, + Some(projection_vector), + if common.has_data_filters || !common.data_filters.is_empty() { + Some(data_filters?) + } else { + None + }, + default_values, + common.session_timezone.as_str(), + common.case_sensitive, + common.return_null_struct_if_all_fields_missing, + common.allow_type_promotion, + common.allow_timestamp_ltz_to_ntz, + self.session_ctx(), + common.encryption_enabled, + common.use_field_id, + common.require_field_ids, + rebase_from_file_metadata, + datetime_rebase_mode_in_read, + int96_rebase_mode_in_read, + )?; + Ok(Arc::new(SparkPlan::new(plan_id, scan, vec![]))) + } + + /// Register the scan's object store and convert its proto file list into DataFusion + /// [`PartitionedFile`]. Shared by the NativeScan and DeltaScan arms; empty partitions + /// yield an empty file list (handled by `build_parquet_scan_plan`). + fn prepare_scan_store_and_files( + &self, + common: &spark_operator::NativeScanCommon, + partition_files: &SparkFilePartition, + ) -> Result<(ObjectStoreUrl, ObjectStoreBackend, Vec), ExecutionError> { + let one_file = match partition_files.partitioned_file.first() { + Some(f) => f.file_path.clone(), + None => { + // Empty partition: no store to resolve; the URL and backend are unused because + // the file group is empty and build_parquet_scan_plan returns EmptyExec. + return Ok(( + ObjectStoreUrl::local_filesystem(), + ObjectStoreBackend::Local, + vec![], + )); + } + }; + let object_store_options: HashMap = common + .object_store_options + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + let (object_store_url, _, object_store_backend) = prepare_object_store_with_configs( + self.session_ctx.runtime_env(), + one_file, + &object_store_options, + )?; + let files = self.get_partitioned_files(partition_files, &object_store_options)?; + Ok((object_store_url, object_store_backend, files)) + } + + /// Parse a scan's serialized default values (for columns missing in older files) into the + /// map consumed by the SchemaMapper. Shared by the NativeScan and DeltaScan arms. + fn parse_default_values( + &self, + common: &spark_operator::NativeScanCommon, + required_schema: &SchemaRef, + ) -> Result>, ExecutionError> { + if common.default_values.len() != common.default_values_indexes.len() { + return Err(GeneralError( + "Scan default values and indexes have different lengths".to_string(), + )); + } + if common.default_values.is_empty() { + return Ok(None); + } + common + .default_values + .iter() + .zip(&common.default_values_indexes) + .map(|(expr, offset)| { + let idx = usize::try_from(*offset) + .map_err(|_| GeneralError(format!("Invalid scan default index {offset}")))?; + let field = required_schema.fields().get(idx).ok_or_else(|| { + GeneralError(format!("Scan default index {idx} is outside schema")) + })?; + let value = self.create_default_value(expr, Arc::clone(required_schema))?; + Ok((Column::new(field.name(), idx), value)) + }) + .collect::, ExecutionError>>() + .map(Some) + } + /// get DataFusion PartitionedFiles from a Spark FilePartition fn get_partitioned_files( &self, @@ -1923,162 +2108,24 @@ impl PhysicalPlanner { .as_ref() .ok_or_else(|| GeneralError("NativeScan missing common data".into()))?; - let data_schema = - convert_spark_types_to_arrow_schema(common.data_schema.as_slice()); - let required_schema: SchemaRef = - convert_spark_types_to_arrow_schema(common.required_schema.as_slice()); - let partition_schema: SchemaRef = - convert_spark_types_to_arrow_schema(common.partition_schema.as_slice()); - let projection_vector: Vec = common - .projection_vector - .iter() - .map(|offset| *offset as usize) - .collect(); - let partition_files = scan .file_partition .as_ref() .ok_or_else(|| GeneralError("NativeScan missing file_partition".into()))?; - // Check if this partition has any files (bucketed scan with bucket pruning may have empty partitions) - if partition_files.partitioned_file.is_empty() { - // Match the nonempty scan's projected table schema, including partition - // and constant metadata columns. Fused operators bind against this schema - // even when the bucket has no rows. - let partition_fields: Vec = partition_schema - .fields() - .iter() - .map(|f| { - Arc::new(Field::new(f.name(), f.data_type().clone(), f.is_nullable())) - }) - .collect(); - let table_schema = TableSchema::builder(data_schema) - .with_table_partition_cols(partition_fields) - .build(); - let output_schema = - Arc::new(table_schema.table_schema().project(&projection_vector)?); - let empty_exec = Arc::new(EmptyExec::new(output_schema)); - return Ok(( - vec![], - vec![], - Arc::new(SparkPlan::new(spark_plan.plan_id, empty_exec, vec![])), - )); - } - - // data_filters may reference partition columns and constant metadata columns - // (e.g. `_metadata.file_size`), which the Parquet reader appends after - // required_schema's columns once partition_values are projected into the - // batch. Bind against the combined schema so `Bound` indices resolve - // correctly -- Scala's `exprToProto(filter, scan.output)` - // (CometNativeScan.scala) numbers columns against that same ordering. - let data_filters: Result>, ExecutionError> = - if common.data_filters.is_empty() { - Ok(vec![]) - } else { - let filter_schema: SchemaRef = Arc::new(Schema::new( - required_schema - .fields() - .iter() - .chain(partition_schema.fields().iter()) - .cloned() - .collect::>(), - )); - let float_operands = self.data_filter_float_operands(); - common - .data_filters - .iter() - .map(|expr| { - self.create_data_filter( - expr, - Arc::clone(&filter_schema), - float_operands, - ) - }) - .collect() - }; - - if common.default_values.len() != common.default_values_indexes.len() { - return Err(GeneralError( - "Scan default values and indexes have different lengths".to_string(), - )); - } - let default_values = if common.default_values.is_empty() { - None - } else { - Some( - common - .default_values - .iter() - .zip(&common.default_values_indexes) - .map(|(expr, offset)| { - let idx = usize::try_from(*offset).map_err(|_| { - GeneralError(format!("Invalid scan default index {offset}")) - })?; - let field = required_schema.fields().get(idx).ok_or_else(|| { - GeneralError(format!( - "Scan default index {idx} is outside schema" - )) - })?; - let value = - self.create_default_value(expr, Arc::clone(&required_schema))?; - Ok((Column::new(field.name(), idx), value)) - }) - .collect::, ExecutionError>>()?, - ) - }; - - // Get one file from this partition (we know it's not empty due to early return above) - let one_file = partition_files - .partitioned_file - .first() - .map(|f| f.file_path.clone()) - .expect("partition should have files after empty check"); - - let object_store_options: HashMap = common - .object_store_options - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(); - let (object_store_url, _, object_store_backend) = - prepare_object_store_with_configs( - self.session_ctx.runtime_env(), - one_file, - &object_store_options, - )?; - - // Get files for this partition - let files = self.get_partitioned_files(partition_files, &object_store_options)?; - let file_groups: Vec> = vec![files]; - - let scan = init_datasource_exec( - required_schema, - Some(data_schema), - Some(partition_schema), + let (object_store_url, object_store_backend, files) = + self.prepare_scan_store_and_files(common, partition_files)?; + let scan = self.build_parquet_scan_plan( + spark_plan.plan_id, + common, object_store_url, object_store_backend, - file_groups, - Some(projection_vector), - if common.has_data_filters || !common.data_filters.is_empty() { - Some(data_filters?) - } else { - None - }, - default_values, - common.session_timezone.as_str(), - common.case_sensitive, - common.return_null_struct_if_all_fields_missing, - common.allow_type_promotion, - common.allow_timestamp_ltz_to_ntz, - self.session_ctx(), - common.encryption_enabled, - common.use_field_id, - common.require_field_ids, + files, + false, + "", + "", )?; - Ok(( - vec![], - vec![], - Arc::new(SparkPlan::new(spark_plan.plan_id, scan, vec![])), - )) + Ok((vec![], vec![], scan)) } OpStruct::CsvScan(scan) => { let data_schema = convert_spark_types_to_arrow_schema(scan.data_schema.as_slice()); @@ -2218,6 +2265,12 @@ impl PhysicalPlanner { if let Some(result) = delta_scan::try_plan_contrib_scan(self, spark_plan, contrib) { return result; } + #[cfg(feature = "delta")] + if let Some(result) = + delta_spark_scan::try_plan_contrib_scan(self, spark_plan, contrib) + { + return result; + } #[cfg(feature = "contrib-lance")] if let Some(result) = lance_scan::try_plan_contrib_scan(self, spark_plan, contrib) { return result; @@ -5348,6 +5401,21 @@ mod tests { assert_eq!(field.metadata().get("source"), Some(&"spark".to_string())); } + /// Pack a `DeltaSparkScan` into the generic `ContribScan` envelope exactly as the + /// contrib jar does on the JVM side. + fn delta_spark_envelope(scan: spark_operator::DeltaSparkScan) -> Operator { + use prost::Message; + Operator { + plan_id: 0, + sql_text_pool: vec![], + children: vec![], + op_struct: Some(OpStruct::ContribScan(spark_operator::ContribScan { + type_url: "type.googleapis.com/comet.contrib.delta_spark.DeltaSparkScan".into(), + value: scan.encode_to_vec(), + })), + } + } + fn create_sort_order(index: i32, type_id: i32, descending: bool, nulls_first: bool) -> Expr { Expr { expr_struct: Some(SortOrder(Box::new(spark_expression::SortOrder { @@ -5687,6 +5755,88 @@ mod tests { } } + #[test] + fn delta_scan_errors_without_delta_feature() { + let op = delta_spark_envelope(spark_operator::DeltaSparkScan { + common: None, + delta_common: None, + file_partition: None, + }); + let planner = PhysicalPlanner::default(); + let err = planner.create_plan(&op, &mut vec![], 1).unwrap_err(); + let msg = format!("{err}"); + #[cfg(not(feature = "delta"))] + assert!( + msg.contains("built without a contrib that handles it"), + "expected mismatched-build error, got: {msg}" + ); + #[cfg(feature = "delta")] + assert!( + msg.contains("missing common data"), + "expected missing-common-data error for an empty DeltaSparkScan, got: {msg}" + ); + } + + #[cfg(feature = "delta")] + fn delta_scan_op(files: Vec) -> Operator { + delta_spark_envelope(spark_operator::DeltaSparkScan { + common: Some(Default::default()), + delta_common: None, + file_partition: Some(spark_operator::DeltaSparkFilePartition { + partitioned_file: files, + }), + }) + } + + #[cfg(feature = "delta")] + #[test] + fn delta_scan_rejects_dv_without_source() { + let op = delta_scan_op(vec![spark_operator::DeltaSparkPartitionedFile { + file: Some(spark_operator::SparkPartitionedFile { + file_path: "file:///tmp/f.parquet".into(), + start: 0, + length: 0, + file_size: 0, + partition_values: vec![], + }), + dv: Some(spark_operator::DeltaSparkDvDescriptor { + storage_type: "u".into(), + absolute_path: None, + inline_data: None, + offset: Some(1), + size_in_bytes: 1, + cardinality: 1, + }), + // (file_path carries a scheme because store resolution now precedes + // the DV handling) + }]); + let err = PhysicalPlanner::default() + .create_plan(&op, &mut vec![], 1) + .unwrap_err(); + let msg = format!("{err}"); + assert!( + msg.contains("neither inline data nor a path"), + "expected malformed-descriptor error, got: {msg}" + ); + } + + #[cfg(feature = "delta")] + #[test] + fn delta_scan_rejects_missing_inner_file() { + let op = delta_scan_op(vec![spark_operator::DeltaSparkPartitionedFile { + file: None, + dv: None, + }]); + let err = PhysicalPlanner::default() + .create_plan(&op, &mut vec![], 1) + .unwrap_err(); + let msg = format!("{err}"); + assert!( + msg.contains("missing inner file"), + "expected missing-inner-file error, got: {msg}" + ); + } + #[test] fn test_unpack_dictionary_primitive() { let op_scan = Operator { diff --git a/native/core/src/execution/planner/delta_spark_scan.rs b/native/core/src/execution/planner/delta_spark_scan.rs new file mode 100644 index 00000000000..c4e6d08d950 --- /dev/null +++ b/native/core/src/execution/planner/delta_spark_scan.rs @@ -0,0 +1,928 @@ +// 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. + +//! JVM-planned Delta handler for the generic `OpStruct::ContribScan` dispatcher, feature-gated +//! behind `delta`. +//! +//! delta-spark has already done log replay, snapshot resolution, and partition pruning by the +//! time the scan reaches Comet, so the envelope carries a concrete file list (plus deletion +//! vector descriptors) and the read path reuses the exact same shared parquet scan builder as +//! `NativeScan` -- inheriting row-group stats pruning, page-index pruning, and filter pushdown. +//! Sibling of the kernel-planned handler in `delta_scan.rs`; the two claim different +//! `type_url`s within the same `ContribScan` envelope. + +use std::collections::HashMap; +use std::sync::Arc; + +use datafusion::execution::object_store::ObjectStoreUrl; +use datafusion::execution::runtime_env::RuntimeEnv; +use object_store::path::Path; +use object_store::ObjectStore; +use url::Url; + +use datafusion_comet_proto::spark_operator::{ + ContribScan, DeltaSparkScan, Operator, SparkFilePartition, SparkPartitionedFile, +}; +use prost::Message; + +use crate::execution::operators::ExecutionError; +use crate::execution::operators::ExecutionError::GeneralError; +use crate::execution::planner::PhysicalPlanner; +use crate::execution::planner::PlanCreationResult; +use crate::parquet::objectstore::s3_blob_fs_support::normalize_object_store_url; +use crate::parquet::parquet_support::{ + hash_object_store_configs, object_store_registration_url, object_store_url_key, + prepare_object_store_with_config_hash, +}; + +/// Type name the JVM-planned Delta contrib claims within the `ContribScan` envelope. The +/// contrib jar packs a `DeltaSparkScan` with a `type_url` of +/// `type.googleapis.com/comet.contrib.delta_spark.DeltaSparkScan`; dispatch keys on the +/// contrib-owned suffix, same convention as the kernel path's `delta_scan.rs`. +const DELTA_SPARK_SCAN_TYPE_NAME: &str = "comet.contrib.delta_spark.DeltaSparkScan"; + +/// Contrib entry point for the `OpStruct::ContribScan` dispatcher. Returns `Some(result)` when +/// the envelope carries a JVM-planned Delta scan, or `None` when the `type_url` belongs to some +/// other contrib. +pub(crate) fn try_plan_contrib_scan( + planner: &PhysicalPlanner, + spark_plan: &Operator, + contrib: &ContribScan, +) -> Option { + if !contrib.type_url.ends_with(DELTA_SPARK_SCAN_TYPE_NAME) { + return None; + } + Some( + DeltaSparkScan::decode(contrib.value.as_slice()) + .map_err(|e| { + GeneralError(format!( + "Failed to decode DeltaSparkScan from contrib_scan: {e}" + )) + }) + .and_then(|scan| plan_delta_spark_scan(planner, spark_plan, &scan)), + ) +} + +fn plan_delta_spark_scan( + planner: &PhysicalPlanner, + spark_plan: &Operator, + scan: &DeltaSparkScan, +) -> PlanCreationResult { + // Delta data files are plain parquet, so this reuses NativeScan's parquet scan builder + // (row-group and page-index pruning, filter pushdown); only the file list is + // Delta-specific. delta_common.column_mapping_mode is informational: native code never + // reads it, and field-id matching is switched by common.use_field_id. + let common = scan + .common + .as_ref() + .ok_or_else(|| GeneralError("DeltaSparkScan missing common data".into()))?; + + let delta_partition = scan + .file_partition + .as_ref() + .ok_or_else(|| GeneralError("DeltaSparkScan missing file_partition".into()))?; + + let spark_partition = SparkFilePartition { + partitioned_file: delta_partition + .partitioned_file + .iter() + .map(|f| { + f.file.clone().ok_or_else(|| { + GeneralError("DeltaSparkPartitionedFile missing inner file".into()) + }) + }) + .collect::, _>>()?, + }; + + // Defense-in-depth against a stale or bypassed JVM gate: DeltaScanSupport.declineReason + // (multiStoreReason) already declines data files spanning multiple object-store authorities + // at planning time, but prepare_scan_store_and_files below resolves this whole partition's + // ObjectStoreUrl from the FIRST file only and then strips every other file down to its bare + // object-store path -- a file that actually lives under a different authority would + // silently read through the first file's store handle. Checked here rather than inside + // prepare_scan_store_and_files itself, which is shared with plain NativeScan and out of + // scope for this Delta-specific invariant. + check_same_object_store_authority(&spark_partition.partitioned_file)?; + + let (object_store_url, object_store_backend, mut files) = + planner.prepare_scan_store_and_files(common, &spark_partition)?; + + // Translate deletion vectors into per-file ParquetAccessPlans so deleted + // rows are skipped inside the reader (composing, by intersection, with + // page-index pruning). Fetching the bitmaps and footers is async I/O; + // create_plan runs on the JNI task thread outside the tokio context, so + // block_on here is safe and keeps the scan a plain DataSourceExec. + if delta_partition + .partitioned_file + .iter() + .any(|f| f.dv.is_some()) + { + let object_store_options: HashMap = common + .object_store_options + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + let runtime_env = planner.session_ctx.runtime_env(); + // Resolve every store this partition touches here, on the JNI thread outside the + // async DV runtime below: a cold S3 store's own internal block_on calls panic when + // nested inside get_runtime().block_on(..). See attach_access_plans's doc comment. + let mut resolver = + PartitionStoreResolver::new(Arc::clone(&runtime_env), &object_store_options); + + // get_partitioned_files maps 1:1 over the proto file list, so the three sources are + // expected to be index-aligned. `.zip()` truncates silently on a length mismatch instead + // of erroring, so check_zip_lengths asserts the invariant up front rather than trusting + // it implicitly -- a future change to any one of the three builders that drops or adds an + // element would otherwise corrupt file-to-DV pairing without either side noticing. + check_zip_lengths( + files.len(), + spark_partition.partitioned_file.len(), + delta_partition.partitioned_file.len(), + )?; + let mut dv_files: Vec = + Vec::with_capacity(files.len()); + let (mut store_references, mut memo_hits) = (0usize, 0usize); + for ((file, spark_file), delta_file) in files + .into_iter() + .zip(spark_partition.partitioned_file.iter()) + .zip(delta_partition.partitioned_file.iter()) + { + let data = resolver.resolve(spark_file.file_path.clone())?; + store_references += 1; + memo_hits += usize::from(data.memo_hit); + let data_store = data.store; + let dv_store = match delta_file + .dv + .as_ref() + .and_then(|dv| dv.absolute_path.clone()) + { + Some(dv_path) => { + let resolved = resolver.resolve(dv_path)?; + store_references += 1; + memo_hits += usize::from(resolved.memo_hit); + Some((resolved.store, resolved.path)) + } + None => None, + }; + dv_files.push(crate::execution::delta_dv::DvScanFile { + file, + file_path: spark_file.file_path.clone(), + dv: delta_file.dv.clone(), + data_store, + dv_store, + }); + } + + // Most references in a partition share one store, so this stays near the reference count. + log::debug!( + "Delta partition resolved {store_references} store references with {memo_hits} memo hits" + ); + files = crate::execution::jni_api::get_runtime().block_on( + crate::execution::delta_dv::attach_access_plans(runtime_env, dv_files), + )?; + } + + // `true`: Delta data files may predate the table (e.g. converted or imported parquet) or + // be written with LEGACY rebase modes, and only each file's own footer metadata can say so + // -- resolve the datetime calendar-rebase policy per file rather than inheriting + // NativeScan's documented no-rebase behavior (see datetime_rebase.rs). The session read + // modes forwarded in delta_common cover files whose metadata does not decide (converted + // non-Spark parquet); absent delta_common (defensive -- the injector always sets it) + // degrades to empty modes, i.e. the EXCEPTION refuse-ancient posture. + let (datetime_rebase_mode, int96_rebase_mode) = scan + .delta_common + .as_ref() + .map(|c| { + ( + c.datetime_rebase_mode_in_read.as_str(), + c.int96_rebase_mode_in_read.as_str(), + ) + }) + .unwrap_or(("", "")); + let scan = planner.build_parquet_scan_plan( + spark_plan.plan_id, + common, + object_store_url, + object_store_backend, + files, + true, + datetime_rebase_mode, + int96_rebase_mode, + )?; + Ok((vec![], vec![], scan)) +} + +/// (scheme, username, host, port), all normalized so equality means "same object-store +/// authority". Scheme and host are lowercased; username (the URI's userinfo -- e.g. the container +/// in `abfss://container@account/...`) is compared verbatim, since object-store identifiers built +/// from it may be case-sensitive and it is safer to draw more authority distinctions than fewer; +/// port is compared as `Option` so an explicit port never collapses into an absent one. +/// Mirrors `DeltaScanSupport.uriAuthority`'s normalization on the JVM side, which folds scheme, +/// userinfo, host, and port into one lowercased `getAuthority`-derived key -- both sides must +/// treat two URIs as the same authority in exactly the same cases so the JVM-side gate +/// (`multiStoreReason`, which declines) always fires before this native check (which errors) ever +/// would. +type ObjectStoreAuthority = (String, String, String, Option); + +/// Errors unless every file in `files` shares the first file's [`ObjectStoreAuthority`]. The +/// `url` crate does NOT lowercase the host for opaque (non-"special") schemes like +/// `s3a`/`abfss`/`hdfs`, so comparing `url[BeforeHost..AfterPort]` verbatim would treat two +/// spellings of the same bucket (`s3a://Bucket-A/..` vs `s3a://bucket-a/..`) as different +/// authorities and hard-error instead of gracefully declining. See the call site's comment for +/// why this defensive check exists alongside the JVM-side gate. +fn check_same_object_store_authority(files: &[SparkPartitionedFile]) -> Result<(), ExecutionError> { + let mut first: Option<(ObjectStoreAuthority, &str)> = None; + for file in files { + let url = Url::parse(&file.file_path).map_err(|e| { + GeneralError(format!( + "Error parsing URL {}: {e}", + redacted_url_display(&file.file_path) + )) + })?; + let authority: ObjectStoreAuthority = ( + url.scheme().to_ascii_lowercase(), + url.username().to_string(), + url.host_str().unwrap_or("").to_ascii_lowercase(), + url.port(), + ); + match &first { + None => first = Some((authority, file.file_path.as_str())), + Some((first_authority, first_path)) if *first_authority != authority => { + return Err(GeneralError(format!( + "Native Delta scan does not support data files spanning multiple object \ + stores (found {} and {})", + redacted_url_display(first_path), + redacted_url_display(&file.file_path) + ))); + } + Some(_) => {} + } + } + Ok(()) +} + +/// Errors unless `files_len`, `spark_files_len`, and `delta_files_len` all agree. Called before +/// the three-way `.zip()` over the object-store-resolved files, the JVM-planned +/// `SparkPartitionedFile`s, and the Delta-specific per-file deletion-vector descriptors that +/// builds `dv_files` -- `Iterator::zip` stops at the shortest sequence with no error, so any +/// future change to one of the three independently-built sources that adds or drops an element +/// would otherwise silently mis-pair a data file with the wrong (or a missing) deletion vector +/// instead of failing loudly. +fn check_zip_lengths( + files_len: usize, + spark_files_len: usize, + delta_files_len: usize, +) -> Result<(), ExecutionError> { + if files_len == spark_files_len && spark_files_len == delta_files_len { + return Ok(()); + } + Err(GeneralError(format!( + "Native Delta scan found mismatched file-list lengths while attaching deletion vectors \ + (resolved files: {files_len}, planned files: {spark_files_len}, deletion-vector \ + descriptors: {delta_files_len}); refusing to zip index-aligned sequences of unequal \ + length" + ))) +} + +/// The userinfo component of `url`'s authority (e.g. the container in +/// `abfss://container@account.dfs.core.windows.net/...`), or the empty string when the URL +/// carries none. Never lowercased, mirroring `check_same_object_store_authority`'s own use of +/// `url.username()` above: userinfo is the ONE component `parquet_support.rs`'s `url_key` drops +/// before it becomes the [`ObjectStoreUrl`] two URLs are resolved and cached under, so it must +/// be compared verbatim, not normalized, to detect a real store-identity collision. Mirrors +/// `DeltaScanSupport.uriUserInfo` on the JVM side. +fn url_user_info(url: &Url) -> String { + url.username().to_string() +} + +/// A display form of `url` safe to embed in an error message: userinfo (e.g. the access/secret +/// key pair embedded as `s3a://AKIA...:secret@bucket/...`, or a Delta shallow-clone container +/// name) is replaced with a literal `***`, mirroring `DeltaScanSupport.redactedAuthority` on the +/// JVM side (`scheme://***@host[:port]`). Scheme and host/port are kept verbatim (not +/// lowercased) and the path is kept in full -- userinfo is the only secret-bearing component, +/// and dropping the path would make the two defense-in-depth checks that call this ([` +/// check_same_object_store_authority`] and [`check_store_identity`]) unable to name which file +/// triggered the error. +/// +/// `url` need not be a valid [`Url`] -- every call site formats a `GeneralError` from a URL that +/// may originate from a foreign/bypassed proto producer, including ones a credential-bearing URL +/// can produce by FAILING to parse in the first place (e.g. `s3a://AKIA:secret@bucket:notaport/x` +/// is `Url::parse`-rejected as `InvalidPort`, but still carries userinfo), so this must be total +/// (never panic) AND must still redact on the parse-failure path -- it is exactly the credentials +/// that make a URL unusual enough to fail parsing that most need to never reach a log line. +/// The fallback below is purely textual: it looks for a `://` scheme delimiter and, within the +/// authority segment that follows (up to the next `/`, mirroring where a real URL's authority +/// ends), replaces everything up to and including the LAST `@` with `***@` -- same last-`@` split +/// as the successfully-parsed path and `DeltaScanSupport.redactedAuthority` on the JVM side. A +/// string with no `://` is treated as having no authority at all and its whole text is searched +/// for a trailing userinfo-shaped `...@host` prefix the same way. A string with neither shape +/// (no `@` anywhere before its authority ends) has no evident secret to redact and is returned +/// unchanged. +fn redacted_url_display(url: &str) -> String { + if let Ok(parsed) = Url::parse(url) { + if parsed.username().is_empty() && parsed.password().is_none() { + return url.to_string(); + } + let host_port = match (parsed.host_str(), parsed.port()) { + (Some(host), Some(port)) => format!("{host}:{port}"), + (Some(host), None) => host.to_string(), + (None, _) => String::new(), + }; + let mut redacted = format!("{}://***@{host_port}{}", parsed.scheme(), parsed.path()); + if let Some(query) = parsed.query() { + redacted.push('?'); + redacted.push_str(query); + } + return redacted; + } + + let (scheme_prefix, rest) = match url.find("://") { + Some(scheme_end) => (&url[..scheme_end + 3], &url[scheme_end + 3..]), + None => ("", url), + }; + let authority_len = rest.find('/').unwrap_or(rest.len()); + match rest[..authority_len].rfind('@') { + Some(at) => format!("{scheme_prefix}***@{}", &rest[at + 1..]), + None => url.to_string(), + } +} + +/// A store resolved by [`PartitionStoreResolver::resolve`]: the within-store `path` of the +/// URL, the `store` handle, and whether the resolver's memo already held that store. +struct ResolvedStore { + path: Path, + store: Arc, + memo_hit: bool, +} + +/// Resolves the object store behind every data-file and deletion-vector URL one partition +/// touches (`check_same_object_store_authority` covers the data files; a DV may live under +/// another authority). Memoizes per registration URL so files sharing an authority pay the +/// global cache lock and `RuntimeEnv` registration once, and hosts the store-identity check. +struct PartitionStoreResolver<'a> { + runtime_env: Arc, + options: &'a HashMap, + /// `options` is the same map for every URL this partition resolves, so it is hashed once. + config_hash: u64, + resolved_stores: HashMap>, + /// Per registration URL, the userinfo and raw URL of the first URL that resolved to it, + /// for [`check_store_identity`]. + store_identities: HashMap, +} + +impl<'a> PartitionStoreResolver<'a> { + fn new(runtime_env: Arc, options: &'a HashMap) -> Self { + Self { + runtime_env, + options, + config_hash: hash_object_store_configs(options), + resolved_stores: HashMap::new(), + store_identities: HashMap::new(), + } + } + + fn resolve(&mut self, url: String) -> Result { + let parsed_url = Url::parse(&url).map_err(|e| { + GeneralError(format!( + "Error parsing URL {}: {e}", + redacted_url_display(&url) + )) + })?; + let user_info = url_user_info(&parsed_url); + // The same normalize-then-key steps the shared resolution path applies, so the memo key + // is the registration URL by construction. Re-parses only so the parse error above stays + // redacted. + let normalized = normalize_object_store_url(&url, self.options)?; + let (url_key, _is_hdfs_scheme) = object_store_url_key(&normalized); + let store_url = object_store_registration_url(&normalized, &url_key, self.config_hash)?; + check_store_identity(&store_url, &user_info, &url, &mut self.store_identities)?; + if let Some(store) = self.resolved_stores.get(&store_url) { + let path = Path::from_url_path(normalized.url.path()) + .map_err(|e| GeneralError(e.to_string()))?; + return Ok(ResolvedStore { + path, + store: Arc::clone(store), + memo_hit: true, + }); + } + + // Memo miss: the expensive path (global cache lock, possible store construction, + // runtime_env registration). It registers under `store_url`, which the drift guard in + // parquet_support's tests pins, so the memo entry is read back under the same key. + let (registered, path, _) = prepare_object_store_with_config_hash( + Arc::clone(&self.runtime_env), + url, + self.options, + self.config_hash, + )?; + debug_assert_eq!( + registered, store_url, + "memo key and registration URL must come from the same derivation" + ); + let store = self.runtime_env.object_store(&store_url)?; + self.resolved_stores.insert(store_url, Arc::clone(&store)); + Ok(ResolvedStore { + path, + store, + memo_hit: false, + }) + } +} + +/// Errors when `store_url` was already resolved earlier in this scan under a DIFFERENT +/// `user_info` than the one now being resolved for `url`; otherwise records `(user_info, url)` +/// for `store_url` in `seen` (first resolution wins the recorded userinfo) and returns `Ok`. +/// +/// This is the free-standing half of the residual cross-container DV check, called from +/// [`PartitionStoreResolver::resolve`] -- the ONE place in this scan that sees both data-file +/// AND deletion-vector URLs. `store_url` is the registration URL the store is memoized and +/// registered under; its authority keeps userinfo only for ABFS, where it is the container +/// (see `object_store_authority` in `parquet_support.rs`), so two URLs of another scheme that +/// agree on `store_url` but disagree on `user_info` are exactly the URLs the native side would +/// otherwise silently collapse onto one store handle. A Delta shallow clone across containers +/// on one storage account, data in `source` and a later DELETE writing its deletion vector +/// into `clone`, resolves each container to its own store and passes this check. +/// +/// Deliberately NOT folded into `check_same_object_store_authority` above: that check only ever +/// sees DATA files and hard-errors on ANY authority mismatch, which would incorrectly reject the +/// legitimate cross-bucket DV shape (data in one S3 bucket, its DV in another) -- distinct hosts +/// mean distinct `store_url`s, so this check never even treats them as collision candidates; see +/// `dv_in_different_bucket_is_allowed` below. +fn check_store_identity( + store_url: &ObjectStoreUrl, + user_info: &str, + url: &str, + seen: &mut HashMap, +) -> Result<(), ExecutionError> { + match seen.get(store_url) { + Some((seen_user_info, seen_url)) if seen_user_info != user_info => { + Err(GeneralError(format!( + "Native Delta scan does not support data files and deletion vectors whose \ + stores collide under the native store-identity key (found {} and {})", + redacted_url_display(seen_url), + redacted_url_display(url) + ))) + } + Some(_) => Ok(()), + None => { + seen.insert(store_url.clone(), (user_info.to_string(), url.to_string())); + Ok(()) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn partitioned_file(path: &str) -> SparkPartitionedFile { + SparkPartitionedFile { + file_path: path.to_string(), + start: 0, + length: 0, + file_size: 0, + partition_values: vec![], + } + } + + #[test] + fn same_authority_files_pass() { + let files = vec![ + partitioned_file("s3a://bucket/a/part-0.parquet"), + partitioned_file("s3a://bucket/b/part-1.parquet"), + ]; + assert!(check_same_object_store_authority(&files).is_ok()); + } + + #[test] + fn same_authority_files_pass_regardless_of_host_case() { + // The `url` crate does not lowercase hosts for opaque (non-"special") schemes like + // s3a, so this must be normalized explicitly rather than relying on Url's own + // formatting -- otherwise the same physical bucket recorded with mixed casing would + // pass the JVM gate (which does lowercase) but hard-error here instead. + let files = vec![ + partitioned_file("s3a://Bucket-A/x.parquet"), + partitioned_file("s3a://bucket-a/y.parquet"), + ]; + assert!(check_same_object_store_authority(&files).is_ok()); + } + + #[test] + fn mixed_authority_files_error_names_both() { + let files = vec![ + partitioned_file("s3a://bucket-a/part-0.parquet"), + partitioned_file("s3a://bucket-b/part-1.parquet"), + ]; + let err = check_same_object_store_authority(&files).unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("bucket-a"), + "expected message to name bucket-a: {msg}" + ); + assert!( + msg.contains("bucket-b"), + "expected message to name bucket-b: {msg}" + ); + assert!( + msg.contains("multiple object stores"), + "expected message to explain the failure: {msg}" + ); + } + + #[test] + fn cross_container_abfss_files_error() { + // Same storage account, different containers: the userinfo (container) must be part of + // the authority key, or `abfss://containerA@account/..` and + // `abfss://containerB@account/..` would collapse into the same authority (same host, + // same scheme) and this defense-in-depth check would silently let a cross-container scan + // through instead of erroring. + let files = vec![ + partitioned_file("abfss://containerA@account.dfs.core.windows.net/a/part-0.parquet"), + partitioned_file("abfss://containerB@account.dfs.core.windows.net/b/part-1.parquet"), + ]; + let err = check_same_object_store_authority(&files).unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("multiple object stores"), + "expected message to explain the failure: {msg}" + ); + } + + #[test] + fn same_container_abfss_files_pass() { + let files = vec![ + partitioned_file("abfss://container@account.dfs.core.windows.net/a/part-0.parquet"), + partitioned_file("abfss://container@account.dfs.core.windows.net/b/part-1.parquet"), + ]; + assert!(check_same_object_store_authority(&files).is_ok()); + } + + #[test] + fn distinct_underscore_host_buckets_error() { + // `gs://my_bucket/..` has an underscore reg-name; the `url` crate (unlike Java's `URI`) + // parses it as an opaque host without failing the whole authority, so this check must + // still tell two distinct underscore-bearing buckets apart. + let files = vec![ + partitioned_file("gs://my_bucket/a/part-0.parquet"), + partitioned_file("gs://other_bucket/b/part-1.parquet"), + ]; + let err = check_same_object_store_authority(&files).unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("multiple object stores"), + "expected message to explain the failure: {msg}" + ); + } + + #[test] + fn same_underscore_host_bucket_files_pass() { + let files = vec![ + partitioned_file("gs://my_bucket/a/part-0.parquet"), + partitioned_file("gs://my_bucket/b/part-1.parquet"), + ]; + assert!(check_same_object_store_authority(&files).is_ok()); + } + + #[test] + fn local_paths_pass_regardless_of_directory() { + let files = vec![ + partitioned_file("file:///tmp/a/part-0.parquet"), + partitioned_file("file:///tmp/b/part-1.parquet"), + ]; + assert!(check_same_object_store_authority(&files).is_ok()); + } + + /// Builds the same `(ObjectStoreUrl, userinfo)` pair `PartitionStoreResolver::resolve` + /// computes for a URL, without touching any object-store backend: it runs the resolver's + /// own normalize-then-key steps, so these fixtures collide (or don't) under + /// [`check_store_identity`] the same way the real resolver's calls would. + fn store_url_and_user_info(url_str: &str) -> (ObjectStoreUrl, String) { + let configs = HashMap::new(); + let user_info = url_user_info(&Url::parse(url_str).unwrap()); + let normalized = normalize_object_store_url(url_str, &configs).unwrap(); + let (key, _) = object_store_url_key(&normalized); + let store_url = + object_store_registration_url(&normalized, &key, hash_object_store_configs(&configs)) + .unwrap(); + (store_url, user_info) + } + + /// Resolves `first` then `second` (same authority) and asserts the second is a memo hit on + /// the same store handle; `third`, when given, must be a miss that adds a second entry. + fn assert_memoized_per_authority( + options: &HashMap, + first: &str, + second: &str, + third: Option<&str>, + ) { + let mut resolver = PartitionStoreResolver::new(Arc::new(RuntimeEnv::default()), options); + let a = resolver.resolve(first.to_string()).unwrap(); + assert!(!a.memo_hit, "{first} must miss a fresh memo"); + let b = resolver.resolve(second.to_string()).unwrap(); + assert!(b.memo_hit, "{second} must hit the memo entry {first} made"); + assert_eq!(resolver.resolved_stores.len(), 1); + assert!(Arc::ptr_eq(&a.store, &b.store)); + assert_eq!( + a.path, + Path::from_url_path(Url::parse(first).unwrap().path()).unwrap() + ); + assert_eq!( + b.path, + Path::from_url_path(Url::parse(second).unwrap().path()).unwrap() + ); + if let Some(third) = third { + let c = resolver.resolve(third.to_string()).unwrap(); + assert!(!c.memo_hit, "{third} must miss: different authority"); + assert_eq!(resolver.resolved_stores.len(), 2); + } + } + + #[test] + #[cfg_attr(miri, ignore)] // AWS credential providers and object_store call foreign functions + fn s3_files_sharing_a_bucket_hit_the_store_memo() { + // The memo key must be the isolated registration URL the resolution path returns + // (`s3+comet--native://bucket`), not the physical `s3://bucket`; otherwise every + // file after the first repeats the global cache lock and runtime registration. + let options = HashMap::from([ + ( + "fs.s3a.aws.credentials.provider".to_string(), + "org.apache.hadoop.fs.s3a.AnonymousAWSCredentialsProvider".to_string(), + ), + ( + "fs.s3a.endpoint.region".to_string(), + "us-east-1".to_string(), + ), + ]); + assert_memoized_per_authority( + &options, + "s3://bucket/a.parquet", + "s3://bucket/b.parquet", + Some("s3://other/c.parquet"), + ); + } + + /// A libhdfs-routed scheme keys its memo entry by the `hdfs` registration URL, so two + /// files of one name node share the entry. The store is seeded in the process-wide cache, + /// since the test build has no name node to construct one against. + #[test] + fn hdfs_files_sharing_a_name_node_hit_the_store_memo() { + use crate::parquet::parquet_support::object_store_cache; + use object_store::memory::InMemory; + let options = HashMap::from([("fs.comet.libhdfs.schemes".to_string(), "hdfs".to_string())]); + let cache_key = ( + "hdfs://comet-memo:8020".to_string(), + hash_object_store_configs(&options), + true, + ); + let store: Arc = Arc::new(InMemory::new()); + object_store_cache() + .write() + .unwrap() + .insert(cache_key.clone(), store); + assert_memoized_per_authority( + &options, + "hdfs://comet-memo:8020/a.parquet", + "hdfs://comet-memo:8020/b.parquet", + None, + ); + object_store_cache().write().unwrap().remove(&cache_key); + } + + #[test] + fn local_files_sharing_a_directory_hit_the_store_memo() { + assert_memoized_per_authority( + &HashMap::new(), + "file:///tmp/x/a.parquet", + "file:///tmp/x/b.parquet", + None, + ); + } + + #[test] + fn dv_in_different_container_same_account_gets_its_own_store() { + // Same storage account (same host), different containers (different userinfo): the + // shape a Delta shallow clone across containers produces when data stays in `source` + // but a later DELETE writes its DV into `clone`. The container is part of the ABFS + // store identity, so the DV resolves through its own store instead of colliding. + let mut seen = HashMap::new(); + let data = "abfss://source@account.dfs.core.windows.net/a/part-0.parquet"; + let dv = "abfss://clone@account.dfs.core.windows.net/_delta_log/deletion_vector_x.bin"; + let (data_store_url, data_user_info) = store_url_and_user_info(data); + let (dv_store_url, dv_user_info) = store_url_and_user_info(dv); + assert_ne!( + data_store_url, dv_store_url, + "containers must not share a store" + ); + check_store_identity(&data_store_url, &data_user_info, data, &mut seen).unwrap(); + check_store_identity(&dv_store_url, &dv_user_info, dv, &mut seen).unwrap(); + assert_eq!(seen.len(), 2); + } + + #[test] + fn dv_with_different_userinfo_on_a_collapsing_scheme_errors() { + // Outside ABFS the store identity drops userinfo, so two URLs that differ only there + // would silently share one store handle and must decline, with the userinfo redacted. + let mut seen = HashMap::new(); + let data = "s3://source@bucket/a/part-0.parquet"; + let dv = "s3://clone@bucket/_delta_log/deletion_vector_x.bin"; + let (data_store_url, data_user_info) = store_url_and_user_info(data); + let (dv_store_url, dv_user_info) = store_url_and_user_info(dv); + assert_eq!( + data_store_url, dv_store_url, + "userinfo must not change the store identity" + ); + check_store_identity(&data_store_url, &data_user_info, data, &mut seen).unwrap(); + let err = check_store_identity(&dv_store_url, &dv_user_info, dv, &mut seen).unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("store-identity"), + "expected message to reference the store-identity collision: {msg}" + ); + assert!( + !msg.contains("source@") && !msg.contains("clone@"), + "expected message to redact the userinfo: {msg}" + ); + assert!( + msg.contains("***@bucket"), + "expected message to show a redacted authority: {msg}" + ); + } + + #[test] + fn dv_in_different_bucket_is_allowed() { + // Guards the legitimate MinIO/S3 shape: data in one bucket, its DV in another. + // Distinct hosts mean distinct ObjectStoreUrls, so these must never even look like a + // collision to check_store_identity -- this is exactly the shape + // check_same_object_store_authority alone would be too strict to allow if the + // collision check were folded into it instead of PartitionStoreResolver::resolve. + let mut seen = HashMap::new(); + let data = "s3://comet-delta-a/part-0.parquet"; + let dv = "s3://comet-delta-b/_delta_log/deletion_vector_x.bin"; + let (data_store_url, data_user_info) = store_url_and_user_info(data); + check_store_identity(&data_store_url, &data_user_info, data, &mut seen).unwrap(); + let (dv_store_url, dv_user_info) = store_url_and_user_info(dv); + assert!(check_store_identity(&dv_store_url, &dv_user_info, dv, &mut seen).is_ok()); + } + + #[test] + fn dv_in_same_container_passes() { + let mut seen = HashMap::new(); + let data = "abfss://container@account.dfs.core.windows.net/a/part-0.parquet"; + let dv = "abfss://container@account.dfs.core.windows.net/_delta_log/deletion_vector_x.bin"; + let (data_store_url, data_user_info) = store_url_and_user_info(data); + check_store_identity(&data_store_url, &data_user_info, data, &mut seen).unwrap(); + let (dv_store_url, dv_user_info) = store_url_and_user_info(dv); + assert!(check_store_identity(&dv_store_url, &dv_user_info, dv, &mut seen).is_ok()); + } + + #[test] + fn dv_with_local_paths_passes() { + let mut seen = HashMap::new(); + let data = "file:///tmp/a/part-0.parquet"; + let dv = "file:///tmp/_delta_log/deletion_vector_x.bin"; + let (data_store_url, data_user_info) = store_url_and_user_info(data); + check_store_identity(&data_store_url, &data_user_info, data, &mut seen).unwrap(); + let (dv_store_url, dv_user_info) = store_url_and_user_info(dv); + assert!(check_store_identity(&dv_store_url, &dv_user_info, dv, &mut seen).is_ok()); + } + + #[test] + fn redacted_url_display_leaves_plain_url_unchanged() { + let url = "s3a://bucket/a/part-0.parquet"; + assert_eq!(redacted_url_display(url), url); + } + + #[test] + fn redacted_url_display_redacts_userinfo() { + let url = "s3a://AKIAEXAMPLE:supersecret@bucket/a/part-0.parquet"; + let redacted = redacted_url_display(url); + assert!( + !redacted.contains("AKIAEXAMPLE") && !redacted.contains("supersecret"), + "expected credentials to be redacted: {redacted}" + ); + assert!( + redacted.contains("bucket"), + "expected host to remain visible: {redacted}" + ); + assert_eq!(redacted, "s3a://***@bucket/a/part-0.parquet"); + } + + #[test] + fn redacted_url_display_redacts_multi_at_password_fully() { + // The '@' inside the password must not be mistaken for the userinfo/host delimiter -- + // the LAST '@' in the authority is the real delimiter, same as the JVM's + // `redactedAuthority` split. + let url = "s3a://user:p@ss@bucket/k"; + let redacted = redacted_url_display(url); + assert!( + !redacted.contains("user") && !redacted.contains("p@ss"), + "expected the entire userinfo, including the embedded '@', to be redacted: {redacted}" + ); + assert_eq!(redacted, "s3a://***@bucket/k"); + } + + #[test] + fn redacted_url_display_is_total_for_non_url_input() { + // Not a valid URL and has no authority-like userinfo prefix before its first '/' -- + // must return unchanged rather than panic. + let input = "not a url at all"; + assert_eq!(redacted_url_display(input), input); + + // Not a valid URL (no scheme, so `Url::parse` rejects it as relative), but does have a + // userinfo-shaped prefix before its first '/' -- must still redact it rather than leak + // it verbatim. + let input = "secret@host/path"; + let redacted = redacted_url_display(input); + assert!( + !redacted.contains("secret"), + "expected the userinfo-shaped prefix to be redacted: {redacted}" + ); + assert_eq!(redacted, "***@host/path"); + } + + #[test] + fn redacted_url_display_redacts_credentials_from_a_scheme_prefixed_url_that_fails_to_parse() { + // Invalid port -- `url::Url::parse` rejects this outright (InvalidPort), so this never + // reaches the successfully-parsed branch above; it must still be caught by the fallback, + // which must recognize the `scheme://` prefix so it doesn't stop at the FIRST '/' in + // that prefix (a bug that would leave userinfo un-redacted for exactly this shape). + let url = "s3a://AKIA:secret@bucket:notaport/path"; + assert!(Url::parse(url).is_err(), "fixture must fail to parse"); + let redacted = redacted_url_display(url); + assert!( + !redacted.contains("AKIA") && !redacted.contains("secret"), + "expected credentials to be redacted: {redacted}" + ); + assert_eq!(redacted, "s3a://***@bucket:notaport/path"); + } + + #[test] + fn zip_lengths_agreeing_pass() { + assert!(check_zip_lengths(3, 3, 3).is_ok()); + assert!(check_zip_lengths(0, 0, 0).is_ok()); + } + + #[test] + fn zip_lengths_mismatch_names_all_three_lengths() { + // Every producer of these three sequences (get_partitioned_files, the + // spark_partition.partitioned_file map, and the raw delta_partition.partitioned_file + // list) currently guarantees 1:1 length agreement on every success path -- this can't be + // reached today through the public ContribScan entry point without a code change + // upstream of this check. It's exercised directly here as defense-in-depth against a + // future regression in one of those producers. + let err = check_zip_lengths(2, 3, 3).unwrap_err(); + let msg = err.to_string(); + assert!(msg.contains("resolved files: 2"), "message was: {msg}"); + assert!(msg.contains("planned files: 3"), "message was: {msg}"); + assert!( + msg.contains("deletion-vector descriptors: 3"), + "message was: {msg}" + ); + } + + #[test] + fn zip_lengths_mismatch_on_delta_files_only() { + let err = check_zip_lengths(4, 4, 5).unwrap_err(); + let msg = err.to_string(); + assert!(msg.contains("resolved files: 4"), "message was: {msg}"); + assert!(msg.contains("planned files: 4"), "message was: {msg}"); + assert!( + msg.contains("deletion-vector descriptors: 5"), + "message was: {msg}" + ); + } + + #[test] + fn parse_error_on_credential_bearing_url_redacts_the_error_message() { + // Regression: a credential-bearing URL that FAILS `Url::parse` (bad port here) must + // still produce an error whose message omits the secret -- this exercises the actual + // `check_same_object_store_authority` error path, not just the helper in isolation. + let files = vec![partitioned_file( + "s3a://AKIA:supersecret@bucket:notaport/part-0.parquet", + )]; + let err = check_same_object_store_authority(&files).unwrap_err(); + let msg = err.to_string(); + assert!( + !msg.contains("AKIA") && !msg.contains("supersecret"), + "expected the parse-error message to redact credentials: {msg}" + ); + assert!( + msg.contains("***@bucket"), + "expected the parse-error message to still name the redacted host: {msg}" + ); + } +} diff --git a/native/core/src/parquet/datetime_rebase.rs b/native/core/src/parquet/datetime_rebase.rs new file mode 100644 index 00000000000..0dd405d27fc --- /dev/null +++ b/native/core/src/parquet/datetime_rebase.rs @@ -0,0 +1,3652 @@ +// 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. + +//! Per-file datetime calendar-rebase handling for the parquet scan. +//! +//! Spark 2.4 and earlier wrote dates and timestamps in the hybrid Julian + Gregorian calendar; +//! Spark 3.0+ uses the proleptic Gregorian calendar and records the calendar policy of every +//! file it writes in the parquet footer's key-value metadata (`org.apache.spark.version`, +//! `org.apache.spark.legacyDateTime`, `org.apache.spark.legacyINT96`, +//! `org.apache.spark.timeZone`). Spark's reader resolves the rebase policy from EACH FILE's +//! writer metadata (`DataSourceUtils.datetimeRebaseSpec` / `int96RebaseSpec`) -- the session's +//! `spark.sql.parquet.datetimeRebaseModeInRead` conf only applies to files whose metadata does +//! not decide the policy on its own -- so a reader that ignores the metadata silently returns +//! values shifted by up to ten days for dates before 1582-10-15 (e.g. `1500-01-01` reads as +//! `1500-01-10`). +//! +//! This module mirrors that per-file resolution: [`resolve_file_rebase_policies`] computes the +//! date / INT64-timestamp / INT96-timestamp policies from a file's arrow schema metadata (the +//! parquet key-value pairs survive the parquet -> arrow schema conversion), and +//! [`wrap_datetime_rebase`] wraps the per-file rewritten expressions' column references in a +//! [`SparkDatetimeRebaseExpr`] that rebases values exactly where that is possible without the +//! JVM's historical timezone tables (dates always; timestamps for a fixed UTC writer zone) and +//! refuses -- rather than silently corrupting -- ancient values it cannot rebase. Nested +//! columns are rebuilt leaf by leaf (struct / list / map / fixed-size list / dictionary), each +//! leaf under its own policy, with nulls and offsets preserved. Modern values are always the +//! identity under every policy: from 1582-10-15 onward for dates, and from +//! [`LAST_SWITCH_JULIAN_TS_SECONDS`] (1900-01-01T00:00:00Z, Spark's +//! `RebaseDateTime.lastSwitchJulianTs`) onward for timestamps. +//! +//! Spark applies `datetimeRebaseSpec` to INT64 `TIMESTAMP_MICROS` / `TIMESTAMP_MILLIS` columns +//! and `int96RebaseSpec` to INT96 columns. The two physical types are indistinguishable in the +//! arrow schema DataFusion hands the expression adapter (both surface as `Timestamp(us, "UTC")` +//! after INT96 coercion), so Comet's parquet reader factory stamps the file's INT96 leaf +//! ordinals -- taken from the parquet footer's own `SchemaDescriptor` -- into the key-value +//! metadata under [`INT96_LEAVES_METADATA_KEY`] before the arrow schema is derived (see +//! [`stamp_int96_leaves`] and `eager_page_index_reader_factory.rs`), and the adapter attributes +//! every timestamp leaf to its spec from that stamp. Without a stamp, the two specs are merged: +//! agreement decides, disagreement degrades to [`RebasePolicy::CheckAncient`]. +//! +//! The wrapper sits BENEATH the schema adapter's nested narrowing (the struct -> struct convert +//! that keeps only the requested children), which is what keeps those ordinals physical -- but +//! it means the wrapper sees every physical child, requested or not. Spark only ever decodes +//! the requested nested schema, so [`FileRebasePolicies::restrict_to_requested`] marks the +//! physical leaves the narrowing drops as the identity: an unrequested ancient `s.ts` never +//! blocks `select s.d`, exactly as in Spark. +//! +//! The same pairing decides a timestamp leaf's policy by the type the query READS it as, since +//! `ParquetVectorUpdaterFactory.getUpdater` keys on the requested Spark type, not the parquet +//! annotation: a leaf read as `TIMESTAMP_NTZ` never rebases (INT96 or INT64; Spark 4.x's +//! `BinaryToSQLTimestampUpdater` / `LongUpdater` consult no mode, and Spark 3.x refuses the +//! INT96 and adjusted-INT64 pairings outright, which Comet's `allow_timestamp_ltz_to_ntz` gate +//! reproduces); a leaf read as `TIMESTAMP` rebases under the datetime spec even when the file +//! declares it `isAdjustedToUTC=false` (`isTimestampTypeMatched` checks the unit only); and a +//! `DATE` leaf keeps the date policy whether read as `DATE` or, on Spark 4.x, as +//! `TIMESTAMP_NTZ` (`DateToTimestampNTZWithRebaseUpdater`). A date or timestamp leaf read as a +//! number (a `DATE` as `INT`, a `TIMESTAMP` of any unit as `BIGINT`, e.g. NANOS under +//! `spark.sql.legacy.parquet.nanosAsLong`) never rebases: the numeric updaters come first in +//! `getUpdater` and return the stored value. Conversely, an `INT32` leaf without a `DATE` +//! annotation read as `DATE` takes the date policy: Spark's INT32 arm picks +//! `IntegerWithRebaseUpdater` from the requested `DateType` alone, so the wrapper rebases the +//! stored days beneath the schema adapter's `Int32 -> Date32` cast. +//! +//! Currently only enabled by the Delta scan arms via +//! `SparkParquetOptions::rebase_from_file_metadata`, which also carries the session read modes +//! ([`SessionRebaseModes`], forwarded from the JVM) that decide the policy for files without +//! Spark writer metadata; the plain NativeScan keeps its documented no-rebase behavior (see +//! the compatibility guide and issue #5010). + +use std::collections::HashMap; +use std::fmt::{self, Display}; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, AsArray, FixedSizeListArray, GenericListArray, MapArray, OffsetSizeTrait, + PrimitiveArray, RecordBatch, StructArray, +}; +use arrow::datatypes::{ + ArrowPrimitiveType, ArrowTimestampType, DataType, Date32Type, FieldRef, Int32Type, Schema, + SchemaRef, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, + TimestampNanosecondType, TimestampSecondType, +}; +use arrow::error::ArrowError; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::common::{DataFusionError, Result as DataFusionResult}; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::ColumnarValue; +use parquet::basic::Type as ParquetPhysicalType; +use parquet::file::metadata::{FileMetaData, KeyValue, ParquetMetaData}; +use parquet::schema::types::SchemaDescriptor; + +use super::name_fold::fold_names; +use super::parquet_support::field_id; + +/// Footer key naming the Spark release that wrote the file; absent for non-Spark writers. +const SPARK_VERSION_METADATA_KEY: &str = "org.apache.spark.version"; +/// Present (empty value) when the file's dates and INT64 timestamps were written with +/// `spark.sql.parquet.datetimeRebaseModeInWrite=LEGACY`. +const SPARK_LEGACY_DATETIME_KEY: &str = "org.apache.spark.legacyDateTime"; +/// Present (empty value) when the file's INT96 timestamps were written with +/// `spark.sql.parquet.int96RebaseModeInWrite=LEGACY`. +const SPARK_LEGACY_INT96_KEY: &str = "org.apache.spark.legacyINT96"; +/// The writer session's time zone, stamped alongside either legacy flag. +const SPARK_TIMEZONE_KEY: &str = "org.apache.spark.timeZone"; + +/// Key-value metadata entry Comet's parquet reader factory adds to a file's footer metadata +/// (in memory only, never written back) so the expression adapter can tell INT96 timestamp +/// columns from INT64 ones after both have been coerced to the same arrow type. Value: +/// `":"`, where leaves are the file's +/// primitive columns in `SchemaDescriptor::columns()` order -- the same depth-first order +/// parquet-rs assigns arrow leaves, so an arrow-side depth-first walk lines up with it. The +/// leaf count lets the reader detect a stamp that does not describe the schema it is paired +/// with (see [`Int96Attribution::from_schema`]). +pub(crate) const INT96_LEAVES_METADATA_KEY: &str = "comet.int96_leaf_columns"; + +/// Day of the Gregorian cutover (1582-10-15) as days since the epoch; rebasing is the identity +/// from this day onward. Same value as Spark's `RebaseDateTime.lastSwitchJulianDay`. +const LAST_SWITCH_JULIAN_DAY: i32 = -141427; + +/// Spark's `RebaseDateTime.lastSwitchJulianTs` (and `lastSwitchGregorianTs`) in seconds since +/// the epoch: 1900-01-01T00:00:00Z. Spark derives it as the latest switch instant across every +/// zone in its `julian-gregorian-rebase-micros.json` table (`getLastSwitchTs`, which also +/// asserts the calendars' difference is zero for every zone from then on): most zones ran on +/// local mean time before 1900, so the last instant at which rebasing changes a value in ANY +/// zone is 1900-01-01T00:00:00Z, not the 1582 cutover. `createTimestampRebaseFuncInRead` +/// under `EXCEPTION` throws exactly for `micros < lastSwitchJulianTs` (after converting +/// `TIMESTAMP_MILLIS` to micros), and `rebaseJulianToGregorianMicros` is the identity from it +/// onward in every zone. The value is in seconds so it scales exactly to any timestamp unit. +pub(crate) const LAST_SWITCH_JULIAN_TS_SECONDS: i64 = -2_208_988_800; + +/// The per-century differences between the Julian and proleptic Gregorian calendars, and the +/// Julian-calendar switch days at which each difference starts to apply. Copied verbatim from +/// Spark's `RebaseDateTime.julianGregDiffs` / `julianGregDiffSwitchDay` (which Spark generated +/// from `localRebaseJulianToGregorianDays`); `rebase_julian_to_gregorian_days` must stay +/// value-for-value equal to Spark's `rebaseJulianToGregorianDays`. +const JULIAN_GREG_DIFFS: [i32; 14] = [2, 1, 0, -1, -2, -3, -4, -5, -6, -7, -8, -9, -10, 0]; +const JULIAN_GREG_DIFF_SWITCH_DAY: [i32; 14] = [ + -719164, -682945, -646420, -609895, -536845, -500320, -463795, -390745, -354220, -317695, + -244645, -208120, -171595, -141427, +]; + +/// Proleptic-Gregorian days since 1970-01-01 for a nominal civil date, via Howard Hinnant's +/// `days_from_civil`. `d` may exceed the month's length; the excess rolls into the following +/// month exactly like `LocalDate.of(y, m, 1).plusDays(d - 1)` in Spark's +/// `localRebaseJulianToGregorianDays` (how the non-existent proleptic date `1000-02-29`, +/// valid in the Julian calendar, lands on `1000-03-01`). +fn days_from_civil(y: i64, m: i64, d: i64) -> i64 { + let y = if m <= 2 { y - 1 } else { y }; + let era = y.div_euclid(400); + let yoe = y - era * 400; // [0, 399] + let mp = (m + 9) % 12; // [0, 11], March = 0 + let doy = (153 * mp + 2) / 5 + d - 1; + let doe = yoe * 365 + yoe / 4 - yoe / 100 + doy; + era * 146097 + doe - 719468 +} + +/// Julian-calendar civil date `(year, month, day)` for a day count since 1970-01-01 that labels +/// days in the Julian calendar (astronomical year numbering: 1 BCE is year 0). Standard +/// Julian-day-number conversion (E.G. Richards' algorithm), exact for any day. +fn julian_day_to_civil(days: i64) -> (i64, i64, i64) { + // Integer (noon) Julian Day Number of this civil day: 1970-01-01 is JDN 2440588. + let jdn = days + 2_440_588; + let f = jdn + 1401; + let e = 4 * f + 3; + let g = e.rem_euclid(1461) / 4; + let h = 5 * g + 2; + let day = h.rem_euclid(153) / 5 + 1; + let month = (h / 153 + 2).rem_euclid(12) + 1; + let year = e.div_euclid(1461) - 4716 + (14 - month) / 12; + (year, month, day) +} + +/// Exact port of Spark's `RebaseDateTime.rebaseJulianToGregorianDays`: reinterprets a day count +/// written in the hybrid Julian + Gregorian calendar as the proleptic Gregorian day count of the +/// same nominal civil date. Identity for days from 1582-10-15 onward. Days before the tables' +/// range (before Julian `0001-01-01`) take the calendar-arithmetic path, mirroring Spark's +/// `localRebaseJulianToGregorianDays` fallback. +pub(crate) fn rebase_julian_to_gregorian_days(days: i32) -> i32 { + if days < JULIAN_GREG_DIFF_SWITCH_DAY[0] { + let (y, m, d) = julian_day_to_civil(days as i64); + (days_from_civil(y, m, 1) + (d - 1)) as i32 + } else { + // Spark's rebaseDays: linear search from the most recent switch day. + let mut i = JULIAN_GREG_DIFF_SWITCH_DAY.len(); + loop { + i -= 1; + if i == 0 || days >= JULIAN_GREG_DIFF_SWITCH_DAY[i] { + break; + } + } + days + JULIAN_GREG_DIFFS[i] + } +} + +/// Timezone strings from `org.apache.spark.timeZone` that denote a fixed zero-offset zone in +/// both `java.util.TimeZone` and `java.time`. Only for these is timestamp rebasing the pure +/// nominal-date shift [`SparkDatetimeRebaseExpr::rebase_timestamp_utc`] computes; any other (or +/// absent) zone needs the JVM's historical timezone tables and stays on the +/// refuse-ancient-values path. +const UTC_EQUIVALENT_TIMEZONES: [&str; 6] = ["UTC", "Etc/UTC", "GMT", "Etc/GMT", "Z", "+00:00"]; + +/// How the writer's session time zone (if recorded) affects timestamp rebasing. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) enum WriterTimeZone { + /// A fixed zero-offset zone: rebasing reduces to the exact nominal-date shift. + Utc, + /// Any other zone, or none recorded (pre-3.0 files): ancient values cannot be rebased + /// without the JVM's historical timezone data. + OtherOrUnknown, +} + +/// One session-level datetime rebase read mode (a `LegacyBehaviorPolicy` value of +/// `spark.sql.parquet.datetimeRebaseModeInRead` / `int96RebaseModeInRead`), consulted by +/// [`resolve_file_rebase_policies`] ONLY for files whose footer metadata does not decide the +/// policy on its own -- exactly the `getOrElse` fallback in Spark's +/// `DataSourceUtils.getRebaseSpec`. Files that carry `org.apache.spark.version` ignore these +/// modes entirely, on every Spark version. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] +pub(crate) enum RebaseReadMode { + /// Refuse ancient values (Spark raises `SparkUpgradeException`); maps to + /// [`RebasePolicy::CheckAncient`]. The default mirrors the conservative posture used + /// before the conf was plumbed through (and Spark 3.x's own conf default). + #[default] + Exception, + /// Read values as proleptic Gregorian without rebasing. + Corrected, + /// Rebase from the hybrid Julian + Gregorian calendar. + Legacy, +} + +impl RebaseReadMode { + /// Parses a `LegacyBehaviorPolicy` conf value. `SQLConf` validates and upper-cases the + /// session conf, but a per-relation `datetimeRebaseMode` option arrives verbatim, so the + /// match is case-insensitive. Anything unrecognized -- including the empty string a proto + /// producer that predates the field sends -- falls back to [`RebaseReadMode::Exception`], + /// which refuses ancient values rather than silently corrupting them. + pub(crate) fn from_conf_value(value: &str) -> Self { + match value.to_ascii_uppercase().as_str() { + "CORRECTED" => RebaseReadMode::Corrected, + "LEGACY" => RebaseReadMode::Legacy, + _ => RebaseReadMode::Exception, + } + } +} + +/// The session's effective datetime rebase read modes, one per spec class (INT64 +/// dates/timestamps vs INT96 timestamps), forwarded from the JVM at planning time. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] +pub(crate) struct SessionRebaseModes { + /// `spark.sql.parquet.datetimeRebaseModeInRead` (or the relation's `datetimeRebaseMode`). + pub datetime: RebaseReadMode, + /// `spark.sql.parquet.int96RebaseModeInRead` (or the relation's `int96RebaseMode`). + pub int96: RebaseReadMode, +} + +/// Calendar policy of one file's date or timestamp columns, resolved from writer metadata the +/// same way Spark's `DataSourceUtils.getRebaseSpec` resolves it. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) enum RebasePolicy { + /// Written in the proleptic Gregorian calendar; values pass through untouched. + Corrected, + /// Written in the hybrid Julian + Gregorian calendar; values must be rebased. + Legacy(WriterTimeZone), + /// Policy could not be pinned down (contradictory flags, or a non-Spark writer under the + /// `EXCEPTION` read mode): modern values -- identical under either calendar -- pass, + /// ancient values raise. Mirrors Spark's `EXCEPTION` behavior (`SparkUpgradeException`). + CheckAncient, +} + +/// Which of a file's leaf columns are physically INT96, from the stamp the parquet reader +/// factory adds under [`INT96_LEAVES_METADATA_KEY`]. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) enum Int96Attribution { + /// No stamp, or a stamp whose leaf count does not match the schema it arrived with: the + /// INT64 and INT96 timestamp specs cannot be told apart per column and are merged. + Unknown, + /// Sorted leaf ordinals (depth-first over the file schema's primitive columns) that are + /// INT96; every other timestamp leaf is INT64. + Known(Vec), +} + +impl Int96Attribution { + /// Parses the stamp out of `schema`'s metadata and validates its leaf count against the + /// schema's own depth-first leaf count, so a stamp that does not describe this schema (a + /// crafted footer key, or a cached-metadata mismatch) degrades to [`Self::Unknown`]. + fn from_schema(schema: &Schema) -> Self { + let Some(stamp) = schema.metadata().get(INT96_LEAVES_METADATA_KEY) else { + return Int96Attribution::Unknown; + }; + let Some((count, ordinals)) = stamp.split_once(':') else { + return Int96Attribution::Unknown; + }; + let schema_leaves: usize = schema + .fields() + .iter() + .map(|f| leaf_count(f.data_type())) + .sum(); + if count.parse::().ok() != Some(schema_leaves) { + return Int96Attribution::Unknown; + } + let parsed: Option> = if ordinals.is_empty() { + Some(Vec::new()) + } else { + ordinals + .split(',') + .map(|o| o.parse::().ok().filter(|o| *o < schema_leaves)) + .collect() + }; + match parsed { + Some(mut leaves) => { + leaves.sort_unstable(); + Int96Attribution::Known(leaves) + } + None => Int96Attribution::Unknown, + } + } + + /// `Some(true)` / `Some(false)` when the leaf is known to be INT96 / INT64, `None` when + /// the attribution is unknown. + fn is_int96(&self, leaf: usize) -> Option { + match self { + Int96Attribution::Unknown => None, + Int96Attribution::Known(leaves) => Some(leaves.binary_search(&leaf).is_ok()), + } + } +} + +/// The [`INT96_LEAVES_METADATA_KEY`] value describing `schema`: its leaf count and the +/// ordinals of its INT96 primitive columns. +pub(crate) fn int96_leaf_stamp(schema: &SchemaDescriptor) -> String { + let ordinals: Vec = schema + .columns() + .iter() + .enumerate() + .filter(|(_, column)| column.physical_type() == ParquetPhysicalType::INT96) + .map(|(ordinal, _)| ordinal.to_string()) + .collect(); + format!("{}:{}", schema.num_columns(), ordinals.join(",")) +} + +/// Returns a copy of `metadata` whose key-value metadata carries the [`int96_leaf_stamp`] of +/// its own schema, or `None` when it already does (the common case after the first open of a +/// file, since the caller caches the stamped copy). Any pre-existing entry under the key -- +/// a file cannot legitimately carry one -- is replaced, never trusted. Only the file-level +/// key-value list changes; row groups and page indexes are carried over as-is. The parquet +/// API cannot carry a file decryptor, nor `FileMetaData`'s crate-private encryption fields +/// (encryption algorithm, footer signing key metadata), across this rebuild, so callers must +/// not stamp opens that supply decryption properties -- and the only consumer, the Delta +/// scan, declines every encrypted-parquet configuration before planning, so a parquet +/// modular encryption file never reaches this path with or without those properties. +pub(crate) fn stamp_int96_leaves(metadata: &ParquetMetaData) -> Option { + let file_metadata = metadata.file_metadata(); + let stamp = int96_leaf_stamp(file_metadata.schema_descr()); + let existing = file_metadata + .key_value_metadata() + .and_then(|kvs| kvs.iter().find(|kv| kv.key == INT96_LEAVES_METADATA_KEY)) + .and_then(|kv| kv.value.as_deref()); + if existing == Some(stamp.as_str()) { + return None; + } + let mut key_values: Vec = file_metadata + .key_value_metadata() + .map(|kvs| { + kvs.iter() + .filter(|kv| kv.key != INT96_LEAVES_METADATA_KEY) + .cloned() + .collect() + }) + .unwrap_or_default(); + key_values.push(KeyValue::new(INT96_LEAVES_METADATA_KEY.to_string(), stamp)); + let stamped_file_metadata = FileMetaData::new( + file_metadata.version(), + file_metadata.num_rows(), + file_metadata.created_by().map(str::to_string), + Some(key_values), + file_metadata.schema_descr_ptr(), + file_metadata.column_orders().cloned(), + ); + Some( + ParquetMetaData::new(stamped_file_metadata, metadata.row_groups().to_vec()) + .into_builder() + .set_column_index(metadata.column_index().cloned()) + .set_offset_index(metadata.offset_index().cloned()) + .build(), + ) +} + +/// Per-file rebase policies for the three affected column classes, plus the INT96 +/// attribution that selects between the two timestamp specs per leaf. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct FileRebasePolicies { + /// `DATE` columns, governed by `org.apache.spark.legacyDateTime` alone. + pub date: RebasePolicy, + /// INT64 `TIMESTAMP_MICROS` / `TIMESTAMP_MILLIS` columns, adjusted to UTC or not: the + /// datetime spec (same resolution as `date`), as Spark's `ParquetVectorUpdaterFactory` + /// selects for INT64 read as `TIMESTAMP`. + pub int64_timestamp: RebasePolicy, + /// INT96 columns: the INT96 spec (`org.apache.spark.legacyINT96`, min version 3.1.0). + pub int96_timestamp: RebasePolicy, + /// Which timestamp leaves are INT96. See [`Int96Attribution`]. + pub int96_leaves: Int96Attribution, + /// Sorted depth-first leaf ordinals -- over the physical file schema, the same ordinals + /// `int96_leaves` uses -- that the query does not read: nested children the schema + /// adapter's struct narrowing drops before any value leaves the scan. Spark never decodes + /// them either, so their policy is the identity whatever the file's calendar. Empty until + /// [`Self::restrict_to_requested`] runs (every leaf requested). + pub unrequested_leaves: Vec, + /// Sorted physical leaf ordinals of timezone-carrying timestamps (INT96, or INT64 with + /// `isAdjustedToUTC=true`) the query reads as `TIMESTAMP_NTZ`. Spark decodes those with + /// `BinaryToSQLTimestampUpdater` / `LongUpdater`, which never rebase, so their policy is + /// the identity whatever the file's calendar. Filled by [`Self::restrict_to_requested`]. + pub ntz_requested_leaves: Vec, + /// Sorted physical leaf ordinals of timezone-free INT64 timestamps + /// (`isAdjustedToUTC=false`) the query reads as `TIMESTAMP`. Spark's INT64 branch checks + /// only the unit (`isTimestampTypeMatched`) and hands a `TimestampType` request to + /// `LongWithRebaseUpdater` under the datetime spec, adjusted or not, so these leaves take + /// the same policy as adjusted INT64 leaves. Filled by [`Self::restrict_to_requested`]. + pub ltz_requested_leaves: Vec, + /// Sorted physical leaf ordinals of dates and timestamps the query reads as a number + /// (`INT`, `BIGINT`, ...). Spark's `ParquetVectorUpdaterFactory` picks the numeric updaters + /// before any date or timestamp arm, and none of them rebase, so their policy is the + /// identity whatever the file's calendar. Filled by [`Self::restrict_to_requested`]. + pub numeric_requested_leaves: Vec, + /// Sorted physical leaf ordinals of `Int32` leaves (INT32 without a `DATE` annotation) the + /// query reads as `DATE`. Spark's INT32 arm hands any `DateType` request to + /// `IntegerWithRebaseUpdater` under the datetime spec, whatever the annotation, so these + /// leaves take the date policy. Filled by [`Self::restrict_to_requested`]. + pub date_requested_leaves: Vec, +} + +/// The leaf ordinals [`push_unrequested_leaves`] records while pairing a physical type with +/// the type the query reads it as. Each list is emitted in depth-first order, so it is already +/// sorted for the binary searches in [`FileRebasePolicies`]. +#[derive(Debug, Default)] +struct LeafPairing { + unrequested: Vec, + ntz_requested: Vec, + ltz_requested: Vec, + numeric_requested: Vec, + date_requested: Vec, +} + +impl FileRebasePolicies { + /// True when some policy is not the plain proleptic-Gregorian pass-through, i.e. when the + /// per-column wrap in [`wrap_datetime_rebase`] can install anything at all. + pub(crate) fn any_rebase_needed(&self) -> bool { + self.date != RebasePolicy::Corrected + || self.int64_timestamp != RebasePolicy::Corrected + || self.int96_timestamp != RebasePolicy::Corrected + } + + /// True when the query reads the leaf as a date or timestamp: it is requested, and not + /// as a number. + fn is_requested(&self, leaf: usize) -> bool { + self.unrequested_leaves.binary_search(&leaf).is_err() + && self.numeric_requested_leaves.binary_search(&leaf).is_err() + } + + /// The policy of the `Date32` leaf at depth-first ordinal `leaf`: the file's date policy, + /// or the identity when the query does not read that leaf. + fn date_policy(&self, leaf: usize) -> RebasePolicy { + if self.is_requested(leaf) { + self.date + } else { + RebasePolicy::Corrected + } + } + + /// The policy of the `Int32` leaf at depth-first ordinal `leaf`: the file's date policy + /// when the query reads it as `DATE`, otherwise the identity. + fn int32_policy(&self, leaf: usize) -> RebasePolicy { + if self.date_requested_leaves.binary_search(&leaf).is_ok() { + self.date + } else { + RebasePolicy::Corrected + } + } + + /// The policy of the timezone-carrying timestamp leaf at depth-first ordinal `leaf`: the + /// identity when the query does not read it or reads it as `TIMESTAMP_NTZ` (Spark's NTZ + /// updaters never rebase; on Spark 3.x the pairing is refused before any rebase decision, + /// which Comet's `allow_timestamp_ltz_to_ntz` gate reproduces); otherwise its physical + /// type's spec when the attribution is known, or else the two specs merged -- agreement + /// decides, disagreement degrades to [`RebasePolicy::CheckAncient`], which still passes + /// every modern value and refuses only ancient ones. + fn timestamp_policy(&self, leaf: usize) -> RebasePolicy { + if !self.is_requested(leaf) || self.ntz_requested_leaves.binary_search(&leaf).is_ok() { + return RebasePolicy::Corrected; + } + match self.int96_leaves.is_int96(leaf) { + Some(true) => self.int96_timestamp, + Some(false) => self.int64_timestamp, + None if self.int64_timestamp == self.int96_timestamp => self.int64_timestamp, + None => RebasePolicy::CheckAncient, + } + } + + /// The policy of the timezone-free timestamp leaf at depth-first ordinal `leaf` (INT64 with + /// `isAdjustedToUTC=false`): the identity unless the query reads it as `TIMESTAMP`, which + /// Spark decodes with `LongWithRebaseUpdater` under the datetime spec exactly like an + /// adjusted INT64 leaf. The stamp is still consulted so a leaf it names INT96 follows the + /// INT96 spec; without a stamp the physical type itself proves INT64, so the two specs are + /// not merged. + fn tz_free_timestamp_policy(&self, leaf: usize) -> RebasePolicy { + if !self.is_requested(leaf) || self.ltz_requested_leaves.binary_search(&leaf).is_err() { + return RebasePolicy::Corrected; + } + match self.int96_leaves.is_int96(leaf) { + Some(true) => self.int96_timestamp, + _ => self.int64_timestamp, + } + } + + /// These policies with every physical leaf the query does not read marked the identity, + /// and every timestamp leaf whose requested type differs from its physical one in timezone + /// presence recorded, so [`leaf_policies`] can pick the policy Spark's + /// `ParquetVectorUpdaterFactory.getUpdater` picks for the REQUESTED type. + /// `requested` pairs each top-level field of `physical_schema` (by position) with the type + /// of the logical field the schema adapter narrows it to -- `None` for a column without a + /// logical counterpart, whose leaves are left as they are (no expression reads it anyway). + /// Nested children pair the way the adapter's struct convert selects them (see + /// [`push_unrequested_leaves`]); the INT96 attribution is untouched, since the ordinals + /// stay physical. `requested` is parallel to the schema's fields; should a caller pass a + /// shorter slice, the trailing columns simply keep every leaf (the safe direction). + pub(crate) fn restrict_to_requested( + mut self, + physical_schema: &Schema, + requested: &[Option<&DataType>], + case_sensitive: bool, + use_field_id: bool, + ) -> Self { + debug_assert_eq!(requested.len(), physical_schema.fields().len()); + let matching = FieldMatching { + case_sensitive, + use_field_id, + }; + let mut next_leaf = 0; + let mut pairing = LeafPairing::default(); + for (field, requested) in physical_schema.fields().iter().zip(requested) { + match requested { + Some(logical) => push_unrequested_leaves( + field.data_type(), + logical, + &mut next_leaf, + matching, + &mut pairing, + ), + None => next_leaf += leaf_count(field.data_type()), + } + } + // Emitted in depth-first order, so already sorted for the binary searches. + self.unrequested_leaves = pairing.unrequested; + self.ntz_requested_leaves = pairing.ntz_requested; + self.ltz_requested_leaves = pairing.ltz_requested; + self.numeric_requested_leaves = pairing.numeric_requested; + self.date_requested_leaves = pairing.date_requested; + self + } +} + +/// The field-matching rules of the schema adapter's nested narrowing +/// (`parquet_convert_struct_to_struct`): names fold per `case_sensitive`, and Parquet field ids +/// select fields when `use_field_id` is set. +#[derive(Debug, Clone, Copy)] +struct FieldMatching { + case_sensitive: bool, + use_field_id: bool, +} + +/// Appends to `out.unrequested` the depth-first leaf ordinals of `physical` (counting from +/// `next_leaf`, which advances past every leaf of `physical`) that reading it as `requested` +/// drops, and records in `out.ntz_requested` / `out.ltz_requested` the timestamp leaves whose +/// requested type has the opposite timezone presence (a timezone-carrying leaf read as +/// `TIMESTAMP_NTZ`, a timezone-free leaf read as `TIMESTAMP`); the unit is irrelevant to +/// either, as it is to Spark's `isTimestampTypeMatched`. A date leaf read as `TINYINT`, +/// `SMALLINT`, `INT`, `BIGINT` or `DOUBLE`, and a timestamp leaf read as `BIGINT`, go to +/// `out.numeric_requested`: Spark's INT32 and INT64 numeric updaters for those pairings never +/// rebase. Spark accepts a date as `BIGINT` or `DOUBLE` only on 4.x (3.x refuses the read); +/// the identity is still the right policy for it. An `Int32` leaf read as `DATE` goes to +/// `out.date_requested`. The other arrow types of INT32 storage that Spark's INT32 arm also +/// reads as `DATE` (`INT_8`, `INT_16`, unsigned, `DECIMAL`) fail the scan's cast to `Date32`, +/// and `INT64` has no `DATE` updater, so none of them reaches a rebase. +/// +/// Recurses through exactly the pairings `parquet_convert_array` narrows, and no others: a +/// struct child is dropped only when NO requested child selects it by either rule the struct +/// convert uses -- folded name, or Parquet field id when ids are in play -- and an ambiguous +/// child (several requested children select it) is kept; `List` pairs with `List` by element +/// type, and `Map` with a `Map` of the same key ordering by its entries, positionally. Any +/// other pairing -- a `LargeList` / `FixedSizeList` / dictionary, a map whose ordering +/// differs, or a shape mismatch -- is handed to arrow's cast or passed through whole by the +/// convert, so it keeps every leaf under its physical type's policy. Keeping a superset of +/// what the narrowing reads is always safe (a spurious check at worst); dropping a leaf the +/// narrowing reads would skip its rebase, so every doubt resolves to "requested". Timestamp +/// leaves inside those pass-through shapes are never recorded either, so they keep the +/// physical rule (a spurious check for an NTZ request, no rebase for a `TIMESTAMP` request of +/// a timezone-free leaf or a `DATE` request of an `Int32` leaf); Spark's requested schemas +/// never take those arrow shapes. +fn push_unrequested_leaves( + physical: &DataType, + requested: &DataType, + next_leaf: &mut usize, + matching: FieldMatching, + out: &mut LeafPairing, +) { + match (physical, requested) { + (DataType::Timestamp(_, Some(_)), DataType::Timestamp(_, None)) => { + out.ntz_requested.push(*next_leaf); + *next_leaf += 1; + } + (DataType::Timestamp(_, None), DataType::Timestamp(_, Some(_))) => { + out.ltz_requested.push(*next_leaf); + *next_leaf += 1; + } + ( + DataType::Date32, + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::Float64, + ) + | (DataType::Timestamp(_, _), DataType::Int64) => { + out.numeric_requested.push(*next_leaf); + *next_leaf += 1; + } + (DataType::Int32, DataType::Date32) => { + out.date_requested.push(*next_leaf); + *next_leaf += 1; + } + (DataType::Struct(physical_fields), DataType::Struct(requested_fields)) => { + let names: Vec<&str> = physical_fields + .iter() + .chain(requested_fields.iter()) + .map(|f| f.name().as_str()) + .collect(); + // A fold failure means the names could not be compared at all; keeping every leaf + // requested is the safe superset, the same as the pass-through pairings below. + let Ok(folded) = fold_names(&names, matching.case_sensitive) else { + *next_leaf += leaf_count(physical); + return; + }; + let (physical_folded, requested_folded) = folded.split_at(physical_fields.len()); + for (i, child) in physical_fields.iter().enumerate() { + let child_id = if matching.use_field_id { + field_id(child) + } else { + None + }; + let mut selectors = requested_fields.iter().enumerate().filter(|(j, r)| { + requested_folded[*j] == physical_folded[i] + || (child_id.is_some() && field_id(r) == child_id) + }); + match (selectors.next(), selectors.next()) { + (None, _) => { + let n = leaf_count(child.data_type()); + out.unrequested.extend(*next_leaf..*next_leaf + n); + *next_leaf += n; + } + (Some((_, requested_child)), None) => push_unrequested_leaves( + child.data_type(), + requested_child.data_type(), + next_leaf, + matching, + out, + ), + (Some(_), Some(_)) => *next_leaf += leaf_count(child.data_type()), + } + } + } + (DataType::List(physical_item), DataType::List(requested_item)) => push_unrequested_leaves( + physical_item.data_type(), + requested_item.data_type(), + next_leaf, + matching, + out, + ), + ( + DataType::Map(physical_entries, physical_sorted), + DataType::Map(requested_entries, requested_sorted), + ) if physical_sorted == requested_sorted => { + match (physical_entries.data_type(), requested_entries.data_type()) { + (DataType::Struct(physical_kv), DataType::Struct(requested_kv)) + if physical_kv.len() == requested_kv.len() => + { + for (p, r) in physical_kv.iter().zip(requested_kv.iter()) { + push_unrequested_leaves( + p.data_type(), + r.data_type(), + next_leaf, + matching, + out, + ); + } + } + _ => *next_leaf += leaf_count(physical), + } + } + _ => *next_leaf += leaf_count(physical), + } +} + +/// The writer time zone recorded in `metadata`, classified for timestamp rebasing. Mirrors the +/// `Option(lookupFileMeta(SPARK_TIMEZONE_METADATA_KEY))` lookup Spark's `getRebaseSpec` performs +/// for every LEGACY resolution, conf-fallback included; Spark substitutes the JVM default zone +/// when the key is absent (`RebaseSpec.timeZone`), which is unavailable natively, so an absent or +/// non-UTC zone classifies as [`WriterTimeZone::OtherOrUnknown`] (dates still rebase fully -- +/// the day rebase is zone-free -- while ancient timestamps refuse rather than guess). +fn writer_time_zone(metadata: &HashMap) -> WriterTimeZone { + match metadata.get(SPARK_TIMEZONE_KEY) { + Some(tz) if UTC_EQUIVALENT_TIMEZONES.contains(&tz.as_str()) => WriterTimeZone::Utc, + _ => WriterTimeZone::OtherOrUnknown, + } +} + +/// One spec resolution, mirroring Spark's `DataSourceUtils.getRebaseSpec` exactly: a Spark +/// version below `min_version` (lexicographic comparison, same as the Scala `String.<`) or a +/// present legacy flag means LEGACY; a Spark version at/after `min_version` without the flag +/// means CORRECTED; no Spark version at all falls back to `conf_mode`, the session read conf +/// forwarded from the JVM (`getRebaseSpec`'s `modeByConfig` fallback, its ONLY use of the +/// conf): CORRECTED passes values through, LEGACY rebases (with the writer zone from the +/// file's `org.apache.spark.timeZone` key, same lookup as the metadata-driven LEGACY path), +/// and EXCEPTION refuses ancient values as [`RebasePolicy::CheckAncient`]. +fn resolve_spec( + metadata: &HashMap, + min_version: &str, + legacy_key: &str, + conf_mode: RebaseReadMode, +) -> RebasePolicy { + match metadata.get(SPARK_VERSION_METADATA_KEY) { + None => match conf_mode { + RebaseReadMode::Corrected => RebasePolicy::Corrected, + RebaseReadMode::Legacy => RebasePolicy::Legacy(writer_time_zone(metadata)), + RebaseReadMode::Exception => RebasePolicy::CheckAncient, + }, + Some(version) => { + if version.as_str() < min_version || metadata.contains_key(legacy_key) { + RebasePolicy::Legacy(writer_time_zone(metadata)) + } else { + RebasePolicy::Corrected + } + } + } +} + +/// Resolves the per-file rebase policies from a file's arrow schema: the parquet footer's +/// key-value pairs in its metadata decide the specs (the datetime spec uses min version +/// `3.0.0` and the INT96 spec `3.1.0`, matching `DataSourceUtils.datetimeRebaseSpec` / +/// `int96RebaseSpec`; `session_modes` supplies the per-spec conf fallback for files without +/// Spark writer metadata), and the reader factory's INT96 stamp -- validated against the +/// schema's leaf structure -- attributes each timestamp leaf to its spec. +pub(crate) fn resolve_file_rebase_policies( + physical_file_schema: &Schema, + session_modes: SessionRebaseModes, +) -> FileRebasePolicies { + let metadata = physical_file_schema.metadata(); + let datetime_spec = resolve_spec( + metadata, + "3.0.0", + SPARK_LEGACY_DATETIME_KEY, + session_modes.datetime, + ); + let int96_spec = resolve_spec( + metadata, + "3.1.0", + SPARK_LEGACY_INT96_KEY, + session_modes.int96, + ); + FileRebasePolicies { + date: datetime_spec, + int64_timestamp: datetime_spec, + int96_timestamp: int96_spec, + int96_leaves: Int96Attribution::from_schema(physical_file_schema), + unrequested_leaves: Vec::new(), + ntz_requested_leaves: Vec::new(), + ltz_requested_leaves: Vec::new(), + numeric_requested_leaves: Vec::new(), + date_requested_leaves: Vec::new(), + } +} + +/// Number of primitive leaves `dt` contains in a depth-first walk -- the same count and order +/// parquet-rs uses when it maps the file's `SchemaDescriptor` columns onto the arrow schema, so +/// arrow-side leaf ordinals line up with [`int96_leaf_stamp`]'s. +fn leaf_count(dt: &DataType) -> usize { + match dt { + DataType::Struct(fields) => fields.iter().map(|f| leaf_count(f.data_type())).sum(), + DataType::List(f) + | DataType::LargeList(f) + | DataType::FixedSizeList(f, _) + | DataType::ListView(f) + | DataType::LargeListView(f) + | DataType::Map(f, _) => leaf_count(f.data_type()), + DataType::Dictionary(_, value) => leaf_count(value), + DataType::RunEndEncoded(_, value) => leaf_count(value.data_type()), + DataType::Union(fields, _) => fields.iter().map(|(_, f)| leaf_count(f.data_type())).sum(), + _ => 1, + } +} + +/// Appends the policy of every leaf of `dt`, in depth-first order, to `out`, consuming leaf +/// ordinals from `next_leaf` (exactly [`leaf_count`] of them). Only `Date32`, timestamps and +/// `Int32` leaves read as `DATE` have a policy to apply, and only when the query reads the +/// leaf; a timestamp leaf's policy follows the type the query reads it as (see +/// [`FileRebasePolicies::timestamp_policy`] and +/// [`FileRebasePolicies::tz_free_timestamp_policy`]), and every other leaf is the identity +/// ([`RebasePolicy::Corrected`]). +fn leaf_policies( + dt: &DataType, + next_leaf: &mut usize, + policies: &FileRebasePolicies, + out: &mut Vec, +) { + match dt { + DataType::Date32 => { + out.push(policies.date_policy(*next_leaf)); + *next_leaf += 1; + } + DataType::Int32 => { + out.push(policies.int32_policy(*next_leaf)); + *next_leaf += 1; + } + DataType::Timestamp(_, Some(_)) => { + out.push(policies.timestamp_policy(*next_leaf)); + *next_leaf += 1; + } + DataType::Timestamp(_, None) => { + out.push(policies.tz_free_timestamp_policy(*next_leaf)); + *next_leaf += 1; + } + DataType::Struct(fields) => { + for f in fields { + leaf_policies(f.data_type(), next_leaf, policies, out); + } + } + // Mirrors `leaf_count` variant for variant, so a rebase-affected leaf inside a nested + // type `rebase_array` cannot rebuild (views, run-end, union -- never produced from a + // parquet schema) still gets its real policy and makes `rebase_array` refuse loudly + // instead of being stamped the identity. + DataType::List(f) + | DataType::LargeList(f) + | DataType::FixedSizeList(f, _) + | DataType::ListView(f) + | DataType::LargeListView(f) + | DataType::Map(f, _) => leaf_policies(f.data_type(), next_leaf, policies, out), + DataType::Dictionary(_, value) => leaf_policies(value, next_leaf, policies, out), + DataType::RunEndEncoded(_, value) => { + leaf_policies(value.data_type(), next_leaf, policies, out) + } + DataType::Union(fields, _) => { + for (_, f) in fields.iter() { + leaf_policies(f.data_type(), next_leaf, policies, out); + } + } + _ => { + *next_leaf += 1; + out.push(RebasePolicy::Corrected); + } + } +} + +/// Wraps every column reference in `expr` whose physical file type contains a rebase-affected +/// leaf under a policy that needs handling with a [`SparkDatetimeRebaseExpr`] carrying that +/// column's per-leaf policies, so both the per-file projection and the pushed-down predicate +/// evaluate rebased values. Columns whose leaves are all the identity -- unaffected types, +/// affected types under [`RebasePolicy::Corrected`], or leaves the query does not read (see +/// [`FileRebasePolicies::restrict_to_requested`]) -- pass through unwrapped. (The pruning +/// predicates derived from the wrapped predicate treat the wrapper as an opaque expression and +/// skip pruning on those columns -- conservative, since file-level statistics are in the +/// file's own calendar.) +pub(crate) fn wrap_datetime_rebase( + expr: Arc, + physical_schema: &SchemaRef, + policies: &FileRebasePolicies, +) -> DataFusionResult> { + expr.transform(|e| { + let Some(col) = e.downcast_ref::() else { + return Ok(Transformed::no(e)); + }; + // Missing columns were already replaced with literals; any surviving reference is + // physical-schema-indexed. Out-of-range means a non-file column (defensive): skip. + let Some(field) = physical_schema.fields().get(col.index()) else { + return Ok(Transformed::no(e)); + }; + // This column's first leaf ordinal: the leaves of every preceding top-level field. + let mut next_leaf: usize = physical_schema.fields()[..col.index()] + .iter() + .map(|f| leaf_count(f.data_type())) + .sum(); + let mut column_leaf_policies = Vec::with_capacity(leaf_count(field.data_type())); + leaf_policies( + field.data_type(), + &mut next_leaf, + policies, + &mut column_leaf_policies, + ); + if column_leaf_policies + .iter() + .all(|p| *p == RebasePolicy::Corrected) + { + return Ok(Transformed::no(e)); + } + Ok(Transformed::yes(Arc::new(SparkDatetimeRebaseExpr { + child: e, + field: Arc::clone(field), + leaf_policies: column_leaf_policies, + }) as Arc)) + }) + .map(|t| t.data) +} + +/// Applies a file's calendar-rebase policies to one column: rebases exactly where possible, +/// raises on ancient values it cannot rebase, and passes modern values (the identity under +/// every policy) through untouched. Nested columns are rebuilt leaf by leaf with nulls and +/// offsets preserved. See the module doc for the policy table. +#[derive(Debug, Eq)] +struct SparkDatetimeRebaseExpr { + child: Arc, + /// The physical file field this expression reads (type preserved by the rebase). + field: FieldRef, + /// One policy per primitive leaf of `field`'s type, in depth-first order (a single entry + /// for a flat column). At least one is not [`RebasePolicy::Corrected`]. + leaf_policies: Vec, +} + +impl SparkDatetimeRebaseExpr { + /// The refusal error, as an [`ArrowError`] so `try_unary` closures can raise it directly; + /// it converts into a `DataFusionError` at the `?` in `evaluate`. + fn rebase_error(&self, detail: &str) -> ArrowError { + ArrowError::ComputeError(format!( + "Native scan cannot rebase ancient values in column '{}': the file was written \ + with the legacy (hybrid Julian/Gregorian) calendar, or does not declare which \ + calendar it used, and {detail}. Reading it natively would return silently \ + shifted values; disable the native Delta scan \ + (spark.comet.scan.delta.enabled=false) to let Spark read this table", + self.field.name(), + )) + } + + fn internal_error(&self, detail: impl Display) -> DataFusionError { + DataFusionError::Internal(format!( + "SparkDatetimeRebaseExpr on column '{}': {detail}", + self.field.name() + )) + } + + /// Rebases a timestamp column written at a fixed zero-offset zone: shift the nominal day + /// with the exact date table, keep the time of day. Matches Spark's + /// `rebaseJulianToGregorianMicros` for UTC, where the hybrid calendar's day boundaries sit + /// exactly on multiples of a day and no timezone transition can apply (UTC's last switch + /// instant in Spark's rebase table is the 1582-10-15 cutover itself). + fn rebase_timestamp_utc(&self, v: i64, units_per_day: i64) -> Result { + // Compare in days, not units: the cutover day times a nanosecond day does not fit i64. + let day = v.div_euclid(units_per_day); + if day >= LAST_SWITCH_JULIAN_DAY as i64 { + return Ok(v); + } + let time_of_day = v - day * units_per_day; + let day = i32::try_from(day).map_err(|_| { + self.rebase_error("the value is outside the rebaseable timestamp range") + })?; + let rebased = rebase_julian_to_gregorian_days(day) as i64; + rebased + .checked_mul(units_per_day) + .and_then(|d| d.checked_add(time_of_day)) + .ok_or_else(|| self.rebase_error("the rebased value overflows the timestamp range")) + } + + /// Whether every valid value in `array` is at or after the Gregorian cutover, so no rebase + /// and no rejection applies and the batch can pass through untouched. Null-free arrays take + /// the vectorised minimum; arrays with nulls take one validity-aware pass, which beats the + /// null-aware minimum and keeps the pass-through for a column holding a single null. + fn all_modern(array: &PrimitiveArray, cutover: T::Native) -> bool { + match array.nulls() { + None => arrow::compute::min(array).is_none_or(|min| min >= cutover), + Some(nulls) => array + .values() + .iter() + .zip(nulls.iter()) + .all(|(&v, valid)| !valid || v >= cutover), + } + } + + /// The refuse-ancient-values policy for timestamps: values at or after the cutover pass + /// through, anything earlier is an error naming `detail`. + fn check_ancient_timestamp( + &self, + v: i64, + units_per_second: i64, + detail: &str, + ) -> Result { + if v >= LAST_SWITCH_JULIAN_TS_SECONDS * units_per_second { + Ok(v) + } else { + Err(self.rebase_error(detail)) + } + } + + fn rebase_timestamp_array( + &self, + array: &PrimitiveArray, + policy: RebasePolicy, + units_per_second: i64, + original: &ArrayRef, + ) -> DataFusionResult { + if policy == RebasePolicy::Corrected + || Self::all_modern(array, LAST_SWITCH_JULIAN_TS_SECONDS * units_per_second) + { + return Ok(Arc::clone(original)); + } + let tz = array.timezone().map(Arc::::from); + let rebased: PrimitiveArray = match policy { + RebasePolicy::Corrected => unreachable!("handled above"), + RebasePolicy::Legacy(WriterTimeZone::Utc) => arrow::compute::try_unary(array, |v| { + self.rebase_timestamp_utc(v, units_per_second * 86_400) + })?, + RebasePolicy::Legacy(WriterTimeZone::OtherOrUnknown) => { + arrow::compute::try_unary(array, |v| { + self.check_ancient_timestamp( + v, + units_per_second, + "rebasing timestamps outside a fixed UTC writer zone needs the JVM's \ + historical timezone tables, which are unavailable natively", + ) + })? + } + RebasePolicy::CheckAncient => arrow::compute::try_unary(array, |v| { + self.check_ancient_timestamp( + v, + units_per_second, + "the timestamp's calendar cannot be determined from the file's metadata", + ) + })?, + }; + Ok(Arc::new(rebased.with_timezone_opt(tz))) + } + + /// Rebases days since the epoch: a `Date32` leaf, or an `Int32` leaf read as `DATE` before + /// the schema adapter casts it. + fn rebase_date_array>( + &self, + dates: &PrimitiveArray, + policy: RebasePolicy, + original: &ArrayRef, + ) -> DataFusionResult { + if policy == RebasePolicy::Corrected || Self::all_modern(dates, LAST_SWITCH_JULIAN_DAY) { + return Ok(Arc::clone(original)); + } + let rebased: PrimitiveArray = match policy { + RebasePolicy::Corrected => unreachable!("handled above"), + // The day rebase is a pure calendar reinterpretation, independent of any timezone, + // so every legacy writer zone rebases dates exactly. + RebasePolicy::Legacy(_) => { + arrow::compute::unary::(dates, rebase_julian_to_gregorian_days) + } + RebasePolicy::CheckAncient => { + arrow::compute::try_unary(dates, |v| -> Result { + if v >= LAST_SWITCH_JULIAN_DAY { + Ok(v) + } else { + Err(self.rebase_error( + "the date's calendar cannot be determined from the file's metadata", + )) + } + })? + } + }; + Ok(Arc::new(rebased)) + } + + fn rebase_list( + &self, + list: &GenericListArray, + field: &FieldRef, + cursor: &mut usize, + ) -> DataFusionResult { + let values = self.rebase_array(list.values(), cursor)?; + Ok(Arc::new(GenericListArray::::try_new( + Arc::clone(field), + list.offsets().clone(), + values, + list.nulls().cloned(), + )?)) + } + + /// Applies the leaf policies starting at `cursor` (advanced past every leaf of `array`'s + /// type) to `array`, rebuilding nested arrays around their transformed leaves. Subtrees + /// whose leaves are all the identity are returned as-is without a rebuild. + fn rebase_array(&self, array: &ArrayRef, cursor: &mut usize) -> DataFusionResult { + let dt = array.data_type(); + let n = leaf_count(dt); + let span = self + .leaf_policies + .get(*cursor..*cursor + n) + .ok_or_else(|| { + self.internal_error(format!( + "array of type {dt} does not match the planned leaf layout (leaf {cursor} \ + of {})", + self.leaf_policies.len() + )) + })?; + if span.iter().all(|p| *p == RebasePolicy::Corrected) { + *cursor += n; + return Ok(Arc::clone(array)); + } + match dt { + DataType::Date32 => { + let policy = span[0]; + *cursor += 1; + self.rebase_date_array(array.as_primitive::(), policy, array) + } + DataType::Int32 => { + let policy = span[0]; + *cursor += 1; + self.rebase_date_array(array.as_primitive::(), policy, array) + } + DataType::Timestamp(unit, _) => { + let policy = span[0]; + *cursor += 1; + match unit { + TimeUnit::Second => self.rebase_timestamp_array( + array.as_primitive::(), + policy, + 1, + array, + ), + TimeUnit::Millisecond => self.rebase_timestamp_array( + array.as_primitive::(), + policy, + 1_000, + array, + ), + TimeUnit::Microsecond => self.rebase_timestamp_array( + array.as_primitive::(), + policy, + 1_000_000, + array, + ), + TimeUnit::Nanosecond => self.rebase_timestamp_array( + array.as_primitive::(), + policy, + 1_000_000_000, + array, + ), + } + } + DataType::Struct(fields) => { + let structs = array.as_struct(); + let columns = structs + .columns() + .iter() + .map(|c| self.rebase_array(c, cursor)) + .collect::>>()?; + Ok(Arc::new(StructArray::try_new( + fields.clone(), + columns, + structs.nulls().cloned(), + )?)) + } + DataType::List(field) => self.rebase_list(array.as_list::(), field, cursor), + DataType::LargeList(field) => self.rebase_list(array.as_list::(), field, cursor), + DataType::FixedSizeList(field, size) => { + let list = array.as_fixed_size_list(); + let values = self.rebase_array(list.values(), cursor)?; + Ok(Arc::new(FixedSizeListArray::try_new( + Arc::clone(field), + *size, + values, + list.nulls().cloned(), + )?)) + } + DataType::Map(field, ordered) => { + let map = array.as_map(); + let entries: ArrayRef = Arc::new(map.entries().clone()); + let entries = self.rebase_array(&entries, cursor)?; + Ok(Arc::new(MapArray::try_new( + Arc::clone(field), + map.offsets().clone(), + entries.as_struct().clone(), + map.nulls().cloned(), + *ordered, + )?)) + } + DataType::Dictionary(_, _) => { + let dictionary = array.as_any_dictionary(); + let values = self.rebase_array(dictionary.values(), cursor)?; + Ok(dictionary.with_values(values)) + } + other => Err(self.internal_error(format!( + "cannot rebase values inside unsupported type {other}" + ))), + } + } +} + +impl PartialEq for SparkDatetimeRebaseExpr { + fn eq(&self, other: &Self) -> bool { + self.child.eq(&other.child) + && self.field.eq(&other.field) + && self.leaf_policies == other.leaf_policies + } +} + +impl Hash for SparkDatetimeRebaseExpr { + fn hash(&self, state: &mut H) { + self.child.hash(state); + self.field.hash(state); + self.leaf_policies.hash(state); + } +} + +impl Display for SparkDatetimeRebaseExpr { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "SPARK_DATETIME_REBASE({})", self.field.name()) + } +} + +impl PhysicalExpr for SparkDatetimeRebaseExpr { + fn data_type(&self, _input_schema: &Schema) -> DataFusionResult { + Ok(self.field.data_type().clone()) + } + + fn nullable(&self, _input_schema: &Schema) -> DataFusionResult { + Ok(self.field.is_nullable()) + } + + fn evaluate(&self, batch: &RecordBatch) -> DataFusionResult { + let array = self.child.evaluate(batch)?.into_array(batch.num_rows())?; + let mut cursor = 0; + let rebased = self.rebase_array(&array, &mut cursor)?; + if cursor != self.leaf_policies.len() { + return Err(self.internal_error(format!( + "array of type {} consumed {cursor} of {} planned leaves", + array.data_type(), + self.leaf_policies.len() + ))); + } + Ok(ColumnarValue::Array(rebased)) + } + + fn return_field(&self, _input_schema: &Schema) -> DataFusionResult { + Ok(Arc::clone(&self.field)) + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.child] + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> DataFusionResult> { + assert_eq!(children.len(), 1); + Ok(Arc::new(SparkDatetimeRebaseExpr { + child: children.pop().expect("child"), + field: Arc::clone(&self.field), + leaf_policies: self.leaf_policies.clone(), + })) + } + + fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + Display::fmt(self, f) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Date32Array, Int32Array, Int64Array, ListArray, TimestampMicrosecondArray}; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::Field; + use parquet::schema::parser::parse_message_type; + + /// Julian-calendar civil date -> hybrid day count (the number a legacy writer stores for + /// that nominal date), the inverse of `julian_day_to_civil`. Fliegel-Van Flandern style + /// Julian-calendar JDN formula, exact with euclidean division. + fn julian_civil_to_day(y: i64, m: i64, d: i64) -> i32 { + let a = (14 - m).div_euclid(12); + let y2 = y + 4800 - a; + let m2 = m + 12 * a - 3; + let jdn = d + (153 * m2 + 2).div_euclid(5) + 365 * y2 + y2.div_euclid(4) - 32083; + (jdn - 2_440_588) as i32 + } + + #[test] + fn day_rebase_matches_spark_table_anchors() { + // Julian 0001-01-01 is hybrid day -719164 and proleptic Gregorian 0001-01-01 is day + // -719162 -- the first entry (+2) of Spark's julianGregDiffs table. + assert_eq!(julian_civil_to_day(1, 1, 1), -719164); + assert_eq!(rebase_julian_to_gregorian_days(-719164), -719162); + // Spark's doc example: Julian 1582-01-01 (-141704) rebases to proleptic -141714. + assert_eq!(julian_civil_to_day(1582, 1, 1), -141704); + assert_eq!(rebase_julian_to_gregorian_days(-141704), -141714); + // The last Julian day (1582-10-04) shifts by the full -10; the first Gregorian day + // (1582-10-15, day -141427) and everything after is the identity. + assert_eq!(julian_civil_to_day(1582, 10, 4), -141428); + assert_eq!(rebase_julian_to_gregorian_days(-141428), -141438); + assert_eq!(rebase_julian_to_gregorian_days(-141427), -141427); + assert_eq!(rebase_julian_to_gregorian_days(0), 0); + assert_eq!(rebase_julian_to_gregorian_days(19876), 19876); + } + + #[test] + fn day_rebase_restores_proleptic_1500_01_01() { + // A legacy writer stores proleptic 1500-01-01 as Julian 1500-01-01, which reads + // unrebased as 1500-01-10. Rebasing must restore 1500-01-01. + let stored = julian_civil_to_day(1500, 1, 1); + assert_eq!(stored, days_from_civil(1500, 1, 10) as i32); + assert_eq!( + rebase_julian_to_gregorian_days(stored), + days_from_civil(1500, 1, 1) as i32 + ); + } + + #[test] + fn day_rebase_rolls_julian_only_leap_days_forward() { + // 1500 is a Julian leap year but not a Gregorian one: Julian 1500-02-29 lands on + // proleptic 1500-03-01, mirroring Spark's LocalDate.of(y, m, 1).plusDays trick. + let stored = julian_civil_to_day(1500, 2, 29); + assert_eq!( + rebase_julian_to_gregorian_days(stored), + days_from_civil(1500, 3, 1) as i32 + ); + } + + #[test] + fn day_rebase_falls_back_to_calendar_arithmetic_before_common_era() { + // One day before the table's range: Julian 0000-12-31 -> proleptic 0000-12-31, which + // is days_from_civil(1,1,1) - 1. + let day = julian_civil_to_day(1, 1, 1) - 1; + assert!(day < JULIAN_GREG_DIFF_SWITCH_DAY[0]); + assert_eq!( + rebase_julian_to_gregorian_days(day), + days_from_civil(1, 1, 1) as i32 - 1 + ); + } + + #[test] + fn day_rebase_is_continuous_across_every_table_switch() { + // At each switch day the table's diff takes over from the previous interval; both must + // agree with the calendar-arithmetic ground truth. The hybrid calendar labels days in + // Julian only BEFORE the 1582-10-15 cutover; from the cutover onward it is Gregorian + // and rebasing is the identity. + for &switch in &JULIAN_GREG_DIFF_SWITCH_DAY { + for day in [switch - 1, switch, switch + 1] { + let expected = if day >= LAST_SWITCH_JULIAN_DAY { + day + } else { + let (y, m, d) = julian_day_to_civil(day as i64); + (days_from_civil(y, m, 1) + (d - 1)) as i32 + }; + assert_eq!( + rebase_julian_to_gregorian_days(day), + expected, + "mismatch at hybrid day {day}" + ); + } + } + } + + fn spark_metadata(entries: &[(&str, &str)]) -> HashMap { + entries + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect() + } + + /// A one-column (`Date32`) schema carrying `entries` as its metadata, for spec-resolution + /// tests that only care about the footer key-value pairs. + fn schema_with(entries: &[(&str, &str)]) -> Schema { + Schema::new_with_metadata( + vec![Field::new("d", DataType::Date32, true)], + spark_metadata(entries), + ) + } + + /// The [`SessionRebaseModes`] used by tests that exercise metadata-driven resolution: the + /// default (EXCEPTION, EXCEPTION), matching an unplumbed conf. + fn default_modes() -> SessionRebaseModes { + SessionRebaseModes::default() + } + + fn modes(datetime: RebaseReadMode, int96: RebaseReadMode) -> SessionRebaseModes { + SessionRebaseModes { datetime, int96 } + } + + fn flat_policies( + date: RebasePolicy, + int64_timestamp: RebasePolicy, + int96_timestamp: RebasePolicy, + ) -> FileRebasePolicies { + FileRebasePolicies { + date, + int64_timestamp, + int96_timestamp, + int96_leaves: Int96Attribution::Unknown, + unrequested_leaves: Vec::new(), + ntz_requested_leaves: Vec::new(), + ltz_requested_leaves: Vec::new(), + numeric_requested_leaves: Vec::new(), + date_requested_leaves: Vec::new(), + } + } + + #[test] + fn policies_for_modern_spark_file_without_flags_are_corrected() { + let schema = schema_with(&[(SPARK_VERSION_METADATA_KEY, "3.5.9")]); + let policies = resolve_file_rebase_policies(&schema, default_modes()); + assert_eq!(policies.date, RebasePolicy::Corrected); + assert_eq!(policies.int64_timestamp, RebasePolicy::Corrected); + assert_eq!(policies.int96_timestamp, RebasePolicy::Corrected); + assert!(!policies.any_rebase_needed()); + } + + #[test] + fn policies_for_both_legacy_flags_with_utc_zone_are_legacy_utc() { + let schema = schema_with(&[ + (SPARK_VERSION_METADATA_KEY, "3.5.9"), + (SPARK_LEGACY_DATETIME_KEY, ""), + (SPARK_LEGACY_INT96_KEY, ""), + (SPARK_TIMEZONE_KEY, "UTC"), + ]); + let policies = resolve_file_rebase_policies(&schema, default_modes()); + assert_eq!(policies.date, RebasePolicy::Legacy(WriterTimeZone::Utc)); + assert_eq!( + policies.int64_timestamp, + RebasePolicy::Legacy(WriterTimeZone::Utc) + ); + assert_eq!( + policies.int96_timestamp, + RebasePolicy::Legacy(WriterTimeZone::Utc) + ); + } + + #[test] + fn policies_for_non_utc_writer_zone_mark_the_zone_unusable() { + let schema = schema_with(&[ + (SPARK_VERSION_METADATA_KEY, "3.5.9"), + (SPARK_LEGACY_DATETIME_KEY, ""), + (SPARK_LEGACY_INT96_KEY, ""), + (SPARK_TIMEZONE_KEY, "America/Los_Angeles"), + ]); + let policies = resolve_file_rebase_policies(&schema, default_modes()); + assert_eq!( + policies.date, + RebasePolicy::Legacy(WriterTimeZone::OtherOrUnknown) + ); + assert_eq!( + policies.int64_timestamp, + RebasePolicy::Legacy(WriterTimeZone::OtherOrUnknown) + ); + assert_eq!( + policies.int96_timestamp, + RebasePolicy::Legacy(WriterTimeZone::OtherOrUnknown) + ); + } + + #[test] + fn mixed_flags_without_attribution_degrade_timestamps_to_check_ancient() { + // legacyDateTime present, legacyINT96 absent on a 3.x file, and no INT96 stamp: dates + // are definitely legacy, but a timestamp leaf cannot be attributed to INT64 (legacy) + // vs INT96 (corrected), so the merged policy is CheckAncient. + let schema = schema_with(&[ + (SPARK_VERSION_METADATA_KEY, "3.5.9"), + (SPARK_LEGACY_DATETIME_KEY, ""), + (SPARK_TIMEZONE_KEY, "UTC"), + ]); + let policies = resolve_file_rebase_policies(&schema, default_modes()); + assert_eq!(policies.date, RebasePolicy::Legacy(WriterTimeZone::Utc)); + assert_eq!( + policies.int64_timestamp, + RebasePolicy::Legacy(WriterTimeZone::Utc) + ); + assert_eq!(policies.int96_timestamp, RebasePolicy::Corrected); + assert_eq!(policies.int96_leaves, Int96Attribution::Unknown); + assert_eq!(policies.timestamp_policy(0), RebasePolicy::CheckAncient); + } + + #[test] + fn mixed_flags_with_attribution_follow_each_leafs_physical_type() { + // Same file, but the reader factory stamped which leaves are INT96: leaf 1 is INT96 + // (corrected), leaf 2 is INT64 (legacy UTC). Leaf 0 is the date. + let ts_dt = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())); + let schema = Schema::new_with_metadata( + vec![ + Field::new("d", DataType::Date32, true), + Field::new("ts96", ts_dt.clone(), true), + Field::new("ts", ts_dt, true), + ], + spark_metadata(&[ + (SPARK_VERSION_METADATA_KEY, "3.5.9"), + (SPARK_LEGACY_DATETIME_KEY, ""), + (SPARK_TIMEZONE_KEY, "UTC"), + (INT96_LEAVES_METADATA_KEY, "3:1"), + ]), + ); + let policies = resolve_file_rebase_policies(&schema, default_modes()); + assert_eq!(policies.int96_leaves, Int96Attribution::Known(vec![1])); + assert_eq!(policies.timestamp_policy(1), RebasePolicy::Corrected); + assert_eq!( + policies.timestamp_policy(2), + RebasePolicy::Legacy(WriterTimeZone::Utc) + ); + } + + #[test] + fn int96_attribution_rejects_stamps_that_do_not_describe_the_schema() { + // The stamp's leaf count must equal the schema's depth-first leaf count (2 here: + // s.d and s.ts); anything else -- or unparsable ordinals -- is Unknown. + let ts_dt = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())); + let nested = |stamp: &str| { + Schema::new_with_metadata( + vec![Field::new( + "s", + DataType::Struct( + vec![ + Field::new("d", DataType::Date32, true), + Field::new("ts", ts_dt.clone(), true), + ] + .into(), + ), + true, + )], + spark_metadata(&[(INT96_LEAVES_METADATA_KEY, stamp)]), + ) + }; + assert_eq!( + Int96Attribution::from_schema(&nested("2:1")), + Int96Attribution::Known(vec![1]) + ); + assert_eq!( + Int96Attribution::from_schema(&nested("2:")), + Int96Attribution::Known(vec![]) + ); + for bad in ["3:1", "2:5", "2:x", "garbage", ""] { + assert_eq!( + Int96Attribution::from_schema(&nested(bad)), + Int96Attribution::Unknown, + "stamp {bad:?}" + ); + } + assert_eq!( + Int96Attribution::from_schema(&schema_with(&[])), + Int96Attribution::Unknown + ); + } + + #[test] + fn int96_leaf_stamp_lists_int96_leaf_ordinals_in_depth_first_order() { + let message = "message m { + required int32 id; + optional int96 ts96; + optional group s { + optional int64 ts (TIMESTAMP(MICROS,true)); + optional int96 inner96; + } + optional group l (LIST) { + repeated group list { + optional int96 element; + } + } + }"; + let schema = SchemaDescriptor::new(Arc::new(parse_message_type(message).unwrap())); + assert_eq!(int96_leaf_stamp(&schema), "5:1,3,4"); + + let flat = SchemaDescriptor::new(Arc::new( + parse_message_type("message m { required int32 id; }").unwrap(), + )); + assert_eq!(int96_leaf_stamp(&flat), "1:"); + } + + #[test] + fn stamp_int96_leaves_adds_the_key_once_and_replaces_a_forged_one() { + use parquet::file::properties::WriterProperties; + use parquet::file::reader::{FileReader, SerializedFileReader}; + use parquet::file::writer::SerializedFileWriter; + + let write = |kvs: Option>| -> ParquetMetaData { + let schema = Arc::new( + parse_message_type("message m { required int32 id; optional int96 ts96; }") + .unwrap(), + ); + let mut buffer = Vec::new(); + let props = WriterProperties::builder() + .set_key_value_metadata(kvs) + .build(); + // No row groups: only the footer matters here. + SerializedFileWriter::new(&mut buffer, schema, Arc::new(props)) + .unwrap() + .close() + .unwrap(); + SerializedFileReader::new(bytes::Bytes::from(buffer)) + .unwrap() + .metadata() + .clone() + }; + let stamp_of = |md: &ParquetMetaData| -> Option { + md.file_metadata() + .key_value_metadata() + .and_then(|kvs| kvs.iter().find(|kv| kv.key == INT96_LEAVES_METADATA_KEY)) + .and_then(|kv| kv.value.clone()) + }; + + let plain = write(Some(vec![KeyValue::new( + SPARK_VERSION_METADATA_KEY.to_string(), + "3.5.9".to_string(), + )])); + let stamped = stamp_int96_leaves(&plain).expect("first stamp rebuilds"); + assert_eq!(stamp_of(&stamped).as_deref(), Some("2:1")); + // The original entries survive next to the stamp; nothing else changed. + assert_eq!( + stamped.file_metadata().key_value_metadata().unwrap().len(), + 2 + ); + assert_eq!(stamped.num_row_groups(), plain.num_row_groups()); + assert_eq!( + stamped.file_metadata().num_rows(), + plain.file_metadata().num_rows() + ); + // Already stamped: no rebuild. + assert!(stamp_int96_leaves(&stamped).is_none()); + + // A file that carries the key itself (it cannot legitimately) is never trusted. + let forged = write(Some(vec![KeyValue::new( + INT96_LEAVES_METADATA_KEY.to_string(), + "2:".to_string(), + )])); + let restamped = stamp_int96_leaves(&forged).expect("forged stamp is replaced"); + assert_eq!(stamp_of(&restamped).as_deref(), Some("2:1")); + assert_eq!( + restamped + .file_metadata() + .key_value_metadata() + .unwrap() + .len(), + 1 + ); + } + + #[test] + fn policies_for_pre_spark3_files_are_legacy_with_unknown_zone() { + // Spark 2.4 wrote the hybrid calendar unconditionally and stamped neither the legacy + // flags nor the writer zone; both specs resolve LEGACY via the version comparison. + let schema = schema_with(&[(SPARK_VERSION_METADATA_KEY, "2.4.8")]); + let policies = resolve_file_rebase_policies(&schema, default_modes()); + assert_eq!( + policies.date, + RebasePolicy::Legacy(WriterTimeZone::OtherOrUnknown) + ); + assert_eq!( + policies.int64_timestamp, + RebasePolicy::Legacy(WriterTimeZone::OtherOrUnknown) + ); + assert_eq!( + policies.int96_timestamp, + RebasePolicy::Legacy(WriterTimeZone::OtherOrUnknown) + ); + } + + #[test] + fn policies_for_int96_min_version_gap_follow_each_spec() { + // A 3.0.x file: datetime spec resolves by flag (absent -> CORRECTED) but the INT96 + // spec's min version is 3.1.0, so 3.0.x is LEGACY for INT96. Without attribution the + // disagreement merges to CheckAncient. + let schema = schema_with(&[(SPARK_VERSION_METADATA_KEY, "3.0.3")]); + let policies = resolve_file_rebase_policies(&schema, default_modes()); + assert_eq!(policies.date, RebasePolicy::Corrected); + assert_eq!(policies.int64_timestamp, RebasePolicy::Corrected); + assert_eq!( + policies.int96_timestamp, + RebasePolicy::Legacy(WriterTimeZone::OtherOrUnknown) + ); + assert_eq!(policies.timestamp_policy(0), RebasePolicy::CheckAncient); + } + + #[test] + fn policies_for_non_spark_files_are_check_ancient_by_default() { + // The default session modes are (EXCEPTION, EXCEPTION): a producer that predates the + // mode fields (empty strings) keeps the conservative refuse-ancient posture. + let policies = resolve_file_rebase_policies(&schema_with(&[]), default_modes()); + assert_eq!(policies.date, RebasePolicy::CheckAncient); + assert_eq!(policies.int64_timestamp, RebasePolicy::CheckAncient); + assert_eq!(policies.int96_timestamp, RebasePolicy::CheckAncient); + } + + #[test] + fn rebase_read_mode_parses_conf_values_and_defaults_to_exception() { + assert_eq!( + RebaseReadMode::from_conf_value("CORRECTED"), + RebaseReadMode::Corrected + ); + assert_eq!( + RebaseReadMode::from_conf_value("LEGACY"), + RebaseReadMode::Legacy + ); + assert_eq!( + RebaseReadMode::from_conf_value("EXCEPTION"), + RebaseReadMode::Exception + ); + // Per-relation options arrive verbatim (SQLConf only upper-cases the session conf). + assert_eq!( + RebaseReadMode::from_conf_value("corrected"), + RebaseReadMode::Corrected + ); + // The proto default (producer predates the field) and anything unrecognized refuse + // ancient values rather than silently corrupting them. + assert_eq!( + RebaseReadMode::from_conf_value(""), + RebaseReadMode::Exception + ); + assert_eq!( + RebaseReadMode::from_conf_value("BOGUS"), + RebaseReadMode::Exception + ); + } + + #[test] + fn non_spark_files_follow_corrected_read_modes() { + // Spark 4.0 defaults both read modes to CORRECTED: a non-Spark file's ancient values + // must read as-is (getRebaseSpec's modeByConfig fallback), not refuse. + let policies = resolve_file_rebase_policies( + &schema_with(&[]), + modes(RebaseReadMode::Corrected, RebaseReadMode::Corrected), + ); + assert_eq!(policies.date, RebasePolicy::Corrected); + assert_eq!(policies.int64_timestamp, RebasePolicy::Corrected); + assert_eq!(policies.int96_timestamp, RebasePolicy::Corrected); + assert!(!policies.any_rebase_needed()); + } + + #[test] + fn non_spark_files_follow_legacy_read_modes() { + // LEGACY conf fallback: Spark rebases with the file's recorded writer zone, or the JVM + // default zone when unrecorded -- unavailable natively, so the zone classifies as + // OtherOrUnknown (dates rebase fully, ancient timestamps refuse). + let policies = resolve_file_rebase_policies( + &schema_with(&[]), + modes(RebaseReadMode::Legacy, RebaseReadMode::Legacy), + ); + assert_eq!( + policies.date, + RebasePolicy::Legacy(WriterTimeZone::OtherOrUnknown) + ); + assert_eq!( + policies.int64_timestamp, + RebasePolicy::Legacy(WriterTimeZone::OtherOrUnknown) + ); + assert_eq!( + policies.int96_timestamp, + RebasePolicy::Legacy(WriterTimeZone::OtherOrUnknown) + ); + + // A recorded UTC-equivalent writer zone upgrades the timestamp path to the exact + // rebase, same as the metadata-driven LEGACY branch (getRebaseSpec looks the timezone + // key up for every LEGACY resolution, conf-fallback included). + let policies = resolve_file_rebase_policies( + &schema_with(&[(SPARK_TIMEZONE_KEY, "UTC")]), + modes(RebaseReadMode::Legacy, RebaseReadMode::Legacy), + ); + assert_eq!(policies.date, RebasePolicy::Legacy(WriterTimeZone::Utc)); + assert_eq!( + policies.int64_timestamp, + RebasePolicy::Legacy(WriterTimeZone::Utc) + ); + assert_eq!( + policies.int96_timestamp, + RebasePolicy::Legacy(WriterTimeZone::Utc) + ); + } + + #[test] + fn non_spark_files_with_mixed_read_modes_resolve_each_spec_independently() { + // datetime CORRECTED + int96 EXCEPTION on a metadata-free file: dates and INT64 + // timestamps follow the datetime spec alone (an ancient INT64 timestamp reads + // verbatim), INT96 leaves follow the INT96 spec. + let ts_dt = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())); + let schema = Schema::new_with_metadata( + vec![ + Field::new("ts", ts_dt.clone(), true), + Field::new("ts96", ts_dt, true), + ], + spark_metadata(&[(INT96_LEAVES_METADATA_KEY, "2:1")]), + ); + let policies = resolve_file_rebase_policies( + &schema, + modes(RebaseReadMode::Corrected, RebaseReadMode::Exception), + ); + assert_eq!(policies.date, RebasePolicy::Corrected); + assert_eq!(policies.timestamp_policy(0), RebasePolicy::Corrected); + assert_eq!(policies.timestamp_policy(1), RebasePolicy::CheckAncient); + + // Without the stamp the disagreeing specs merge to CheckAncient for every leaf. + let policies = resolve_file_rebase_policies( + &schema_with(&[]), + modes(RebaseReadMode::Corrected, RebaseReadMode::Legacy), + ); + assert_eq!(policies.date, RebasePolicy::Corrected); + assert_eq!(policies.timestamp_policy(0), RebasePolicy::CheckAncient); + } + + #[test] + fn spark_files_ignore_the_session_read_modes() { + // getRebaseSpec consults modeByConfig ONLY when org.apache.spark.version is absent: a + // legacy 2.4 file stays LEGACY under CORRECTED read modes, and a modern flag-free file + // stays CORRECTED under LEGACY read modes. + let legacy = schema_with(&[(SPARK_VERSION_METADATA_KEY, "2.4.8")]); + let policies = resolve_file_rebase_policies( + &legacy, + modes(RebaseReadMode::Corrected, RebaseReadMode::Corrected), + ); + assert_eq!( + policies.date, + RebasePolicy::Legacy(WriterTimeZone::OtherOrUnknown) + ); + assert_eq!( + policies.int64_timestamp, + RebasePolicy::Legacy(WriterTimeZone::OtherOrUnknown) + ); + assert_eq!( + policies.int96_timestamp, + RebasePolicy::Legacy(WriterTimeZone::OtherOrUnknown) + ); + + let modern = schema_with(&[(SPARK_VERSION_METADATA_KEY, "3.5.9")]); + let policies = resolve_file_rebase_policies( + &modern, + modes(RebaseReadMode::Legacy, RebaseReadMode::Legacy), + ); + assert_eq!(policies.date, RebasePolicy::Corrected); + assert_eq!(policies.int64_timestamp, RebasePolicy::Corrected); + assert_eq!(policies.int96_timestamp, RebasePolicy::Corrected); + } + + /// A wrapper applying `policy` to every leaf of `field` (the same policy for dates and + /// timestamps alike). + fn rebase_expr(field: Field, policy: RebasePolicy) -> SparkDatetimeRebaseExpr { + let policies = flat_policies(policy, policy, policy); + let mut next_leaf = 0; + let mut leaf_pols = Vec::new(); + leaf_policies(field.data_type(), &mut next_leaf, &policies, &mut leaf_pols); + SparkDatetimeRebaseExpr { + child: Arc::new(Column::new(field.name(), 0)), + field: Arc::new(field), + leaf_policies: leaf_pols, + } + } + + fn eval_on( + expr: &SparkDatetimeRebaseExpr, + array: ArrayRef, + field: Field, + ) -> DataFusionResult { + let schema = Arc::new(Schema::new(vec![field])); + let batch = RecordBatch::try_new(schema, vec![array]).unwrap(); + expr.evaluate(&batch)?.into_array(batch.num_rows()) + } + + #[test] + fn legacy_dates_rebase_and_preserve_nulls() { + let field = Field::new("d", DataType::Date32, true); + let expr = rebase_expr(field.clone(), RebasePolicy::Legacy(WriterTimeZone::Utc)); + let stored = julian_civil_to_day(1500, 1, 1); + let array: ArrayRef = Arc::new(Date32Array::from(vec![Some(stored), None, Some(19876)])); + let rebased = eval_on(&expr, array, field).unwrap(); + let rebased = rebased.as_any().downcast_ref::().unwrap(); + assert_eq!(rebased.value(0), days_from_civil(1500, 1, 1) as i32); + assert!(rebased.is_null(1)); + assert_eq!(rebased.value(2), 19876); + } + + #[test] + fn check_ancient_dates_error_only_when_ancient_values_appear() { + let field = Field::new("d", DataType::Date32, true); + let expr = rebase_expr(field.clone(), RebasePolicy::CheckAncient); + let modern: ArrayRef = Arc::new(Date32Array::from(vec![Some(0), Some(19876), None])); + assert!(eval_on(&expr, modern, field.clone()).is_ok()); + + let ancient: ArrayRef = Arc::new(Date32Array::from(vec![Some(-141428)])); + let err = eval_on(&expr, ancient, field).unwrap_err().to_string(); + assert!(err.contains("rebase"), "unexpected error: {err}"); + assert!(err.contains("'d'"), "unexpected error: {err}"); + } + + #[test] + fn modern_batches_pass_through_without_a_new_buffer() { + // A batch with no valid value before the cutover is returned as the same Arc under every + // policy, with or without nulls; null slots may hold ancient garbage and validity decides. + let field = Field::new("d", DataType::Date32, true); + let dates: ArrayRef = Arc::new(Date32Array::from(vec![Some(0), Some(19876)])); + let masked: ArrayRef = { + let values = vec![0i32, -141428, 19876]; + let nulls = arrow::buffer::NullBuffer::from(vec![true, false, true]); + Arc::new(Date32Array::new(values.into(), Some(nulls))) + }; + let ts_field = Field::new( + "ts", + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + true, + ); + let ts: ArrayRef = Arc::new( + TimestampMicrosecondArray::from(vec![ + Some(1_700_000_000_000_000i64), + Some(1_700_000_000_000_001i64), + ]) + .with_timezone("UTC"), + ); + let ts_masked: ArrayRef = { + let values = vec![ + 1_700_000_000_000_000i64, + -14_000_000_000_000_000i64, + 1_700_000_000_000_001i64, + ]; + let nulls = arrow::buffer::NullBuffer::from(vec![true, false, true]); + Arc::new( + TimestampMicrosecondArray::new(values.into(), Some(nulls)).with_timezone("UTC"), + ) + }; + for policy in [ + RebasePolicy::CheckAncient, + RebasePolicy::Legacy(WriterTimeZone::Utc), + RebasePolicy::Legacy(WriterTimeZone::OtherOrUnknown), + ] { + for input in [&dates, &masked] { + let out = eval_on( + &rebase_expr(field.clone(), policy), + Arc::clone(input), + field.clone(), + ) + .unwrap(); + assert!( + Arc::ptr_eq(&out, input), + "{policy:?} should return the input dates" + ); + } + for input in [&ts, &ts_masked] { + let out = eval_on( + &rebase_expr(ts_field.clone(), policy), + Arc::clone(input), + ts_field.clone(), + ) + .unwrap(); + assert!( + Arc::ptr_eq(&out, input), + "{policy:?} should return the input timestamps" + ); + } + } + } + + #[test] + fn ancient_values_still_reject_and_rebase_after_the_pass_through_check() { + let field = Field::new( + "ts", + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + true, + ); + let ancient: ArrayRef = Arc::new( + TimestampMicrosecondArray::from(vec![ + Some(1_700_000_000_000_000i64), + Some(-14_000_000_000_000_000i64), + ]) + .with_timezone("UTC"), + ); + for policy in [ + RebasePolicy::CheckAncient, + RebasePolicy::Legacy(WriterTimeZone::OtherOrUnknown), + ] { + let err = match eval_on( + &rebase_expr(field.clone(), policy), + Arc::clone(&ancient), + field.clone(), + ) { + Ok(out) => panic!("{policy:?} accepted an ancient value: {out:?}"), + Err(e) => e.to_string(), + }; + assert!( + err.contains("rebase"), + "{policy:?}: unexpected error: {err}" + ); + } + let out = eval_on( + &rebase_expr(field.clone(), RebasePolicy::Legacy(WriterTimeZone::Utc)), + Arc::clone(&ancient), + field, + ) + .unwrap(); + assert!( + !Arc::ptr_eq(&out, &ancient), + "an ancient value needs a rebased buffer" + ); + } + + #[test] + fn legacy_utc_timestamps_rebase_by_nominal_day_shift() { + const MICROS_PER_DAY: i64 = 86_400_000_000; + let dt = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())); + let field = Field::new("ts", dt.clone(), true); + let expr = rebase_expr(field.clone(), RebasePolicy::Legacy(WriterTimeZone::Utc)); + // Julian 1500-01-01T12:34:56.789Z as a legacy writer stores it. + let time_of_day = (12i64 * 3600 + 34 * 60 + 56) * 1_000_000 + 789_000; + let stored = julian_civil_to_day(1500, 1, 1) as i64 * MICROS_PER_DAY + time_of_day; + let modern = 1_700_000_000_000_000i64; + let array: ArrayRef = Arc::new( + TimestampMicrosecondArray::from(vec![Some(stored), None, Some(modern)]) + .with_timezone("UTC"), + ); + let rebased = eval_on(&expr, array, field).unwrap(); + assert_eq!(rebased.data_type(), &dt); + let rebased = rebased + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + rebased.value(0), + days_from_civil(1500, 1, 1) * MICROS_PER_DAY + time_of_day + ); + assert!(rebased.is_null(1)); + assert_eq!(rebased.value(2), modern); + } + + #[test] + fn legacy_utc_timestamps_rebase_in_every_unit_without_overflow() { + // The cutover day times a nanosecond day exceeds i64, so the identity check must + // compare in days. Julian 1500-01-01T00:00:01 in each unit that can hold it rebases + // to proleptic 1500-01-01T00:00:01; the epoch, a modern value and -- for nanoseconds, + // whose i64 range only reaches back to 1677 -- i64::MIN are the identity. + let stored_day = julian_civil_to_day(1500, 1, 1) as i64; + let expected_day = days_from_civil(1500, 1, 1); + for (unit, per_second, holds_ancient) in [ + (TimeUnit::Second, 1i64, true), + (TimeUnit::Millisecond, 1_000, true), + (TimeUnit::Microsecond, 1_000_000, true), + (TimeUnit::Nanosecond, 1_000_000_000, false), + ] { + let per_day = per_second * 86_400; + let expr = rebase_expr(ts_field(unit), RebasePolicy::Legacy(WriterTimeZone::Utc)); + let modern = 1_700_000_000 * per_second; + let (ancient_in, ancient_out) = if holds_ancient { + ( + stored_day * per_day + per_second, + expected_day * per_day + per_second, + ) + } else { + (i64::MIN, i64::MIN) + }; + let input = ts_array(unit, vec![Some(ancient_in), Some(0), Some(modern), None]); + let out = + eval_on(&expr, input, ts_field(unit)).unwrap_or_else(|e| panic!("{unit:?}: {e}")); + let expected = ts_array(unit, vec![Some(ancient_out), Some(0), Some(modern), None]); + assert_eq!(&out, &expected, "{unit:?}"); + } + } + + #[test] + fn legacy_non_utc_timestamps_pass_modern_and_refuse_ancient_values() { + let dt = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())); + let field = Field::new("ts", dt, true); + let expr = rebase_expr( + field.clone(), + RebasePolicy::Legacy(WriterTimeZone::OtherOrUnknown), + ); + let modern: ArrayRef = Arc::new( + TimestampMicrosecondArray::from(vec![Some(0), Some(1_700_000_000_000_000)]) + .with_timezone("UTC"), + ); + assert!(eval_on(&expr, modern, field.clone()).is_ok()); + + let ancient: ArrayRef = Arc::new( + TimestampMicrosecondArray::from(vec![Some( + LAST_SWITCH_JULIAN_TS_SECONDS * 1_000_000 - 1, + )]) + .with_timezone("UTC"), + ); + let err = eval_on(&expr, ancient, field).unwrap_err().to_string(); + assert!(err.contains("rebase"), "unexpected error: {err}"); + assert!(err.contains("timezone tables"), "unexpected error: {err}"); + } + + fn ts_field(unit: TimeUnit) -> Field { + Field::new("ts", DataType::Timestamp(unit, Some("UTC".into())), true) + } + + fn ts_array(unit: TimeUnit, values: Vec>) -> ArrayRef { + let tz: Option> = Some("UTC".into()); + match unit { + TimeUnit::Second => { + Arc::new(PrimitiveArray::::from(values).with_timezone_opt(tz)) + } + TimeUnit::Millisecond => Arc::new( + PrimitiveArray::::from(values).with_timezone_opt(tz), + ), + TimeUnit::Microsecond => Arc::new( + PrimitiveArray::::from(values).with_timezone_opt(tz), + ), + TimeUnit::Nanosecond => Arc::new( + PrimitiveArray::::from(values).with_timezone_opt(tz), + ), + } + } + + #[test] + fn check_ancient_timestamps_reject_only_values_before_1900_in_every_unit() { + // Spark's EXCEPTION read mode (`createTimestampRebaseFuncInRead`) throws only for + // micros < RebaseDateTime.lastSwitchJulianTs (1900-01-01T00:00:00Z, the last instant at + // which rebasing changes a value in ANY zone), after converting MILLIS to micros. A + // timestamp one microsecond before the epoch is well after that and must read. + assert_eq!( + LAST_SWITCH_JULIAN_TS_SECONDS, + days_from_civil(1900, 1, 1) * 86_400 + ); + for (unit, per_second) in [ + (TimeUnit::Second, 1i64), + (TimeUnit::Millisecond, 1_000), + (TimeUnit::Microsecond, 1_000_000), + (TimeUnit::Nanosecond, 1_000_000_000), + ] { + let cutoff = LAST_SWITCH_JULIAN_TS_SECONDS * per_second; + for policy in [ + RebasePolicy::CheckAncient, + RebasePolicy::Legacy(WriterTimeZone::OtherOrUnknown), + ] { + let expr = rebase_expr(ts_field(unit), policy); + let passing = ts_array(unit, vec![Some(-1), Some(cutoff), Some(0), None]); + let out = eval_on(&expr, Arc::clone(&passing), ts_field(unit)) + .unwrap_or_else(|e| panic!("{unit:?} under {policy:?}: {e}")); + assert_eq!(&out, &passing, "{unit:?} under {policy:?}"); + + let failing = ts_array(unit, vec![Some(cutoff - 1)]); + let err = eval_on(&expr, failing, ts_field(unit)) + .unwrap_err() + .to_string(); + assert!(err.contains("rebase"), "{unit:?} under {policy:?}: {err}"); + } + } + } + + #[test] + fn wrap_targets_only_affected_columns() { + let policies = flat_policies( + RebasePolicy::Legacy(WriterTimeZone::Utc), + RebasePolicy::Legacy(WriterTimeZone::Utc), + RebasePolicy::Legacy(WriterTimeZone::Utc), + ); + let schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("i", DataType::Int64, true), + Field::new("d", DataType::Date32, true), + Field::new( + "ntz", + DataType::Timestamp(TimeUnit::Microsecond, None), + true, + ), + ])); + let unaffected = wrap_datetime_rebase( + Arc::new(Column::new("i", 0)) as Arc, + &schema, + &policies, + ) + .unwrap(); + assert!(unaffected.downcast_ref::().is_some()); + + let ntz = wrap_datetime_rebase( + Arc::new(Column::new("ntz", 2)) as Arc, + &schema, + &policies, + ) + .unwrap(); + assert!(ntz.downcast_ref::().is_some()); + + let wrapped = wrap_datetime_rebase( + Arc::new(Column::new("d", 1)) as Arc, + &schema, + &policies, + ) + .unwrap(); + let wrapped = wrapped.downcast_ref::().unwrap(); + assert_eq!( + wrapped.leaf_policies, + vec![RebasePolicy::Legacy(WriterTimeZone::Utc)] + ); + } + + #[test] + fn wrap_attributes_timestamp_leaves_by_ordinal_across_preceding_columns() { + // Leaf ordinals count every leaf of the preceding top-level fields: `s` holds leaves + // 0..3 (i, ts96, ts) and the top-level `ts96` is leaf 3. The stamp marks 1 and 3. + let ts_dt = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())); + let schema: SchemaRef = Arc::new(Schema::new_with_metadata( + vec![ + Field::new( + "s", + DataType::Struct( + vec![ + Field::new("i", DataType::Int64, true), + Field::new("ts96", ts_dt.clone(), true), + Field::new("ts", ts_dt.clone(), true), + ] + .into(), + ), + true, + ), + Field::new("ts96", ts_dt, true), + ], + spark_metadata(&[(INT96_LEAVES_METADATA_KEY, "4:1,3")]), + )); + let policies = resolve_file_rebase_policies( + &schema, + modes(RebaseReadMode::Corrected, RebaseReadMode::Exception), + ); + let s = wrap_datetime_rebase( + Arc::new(Column::new("s", 0)) as Arc, + &schema, + &policies, + ) + .unwrap(); + let s = s.downcast_ref::().unwrap(); + assert_eq!( + s.leaf_policies, + vec![ + RebasePolicy::Corrected, + RebasePolicy::CheckAncient, + RebasePolicy::Corrected + ] + ); + let top = wrap_datetime_rebase( + Arc::new(Column::new("ts96", 1)) as Arc, + &schema, + &policies, + ) + .unwrap(); + let top = top.downcast_ref::().unwrap(); + assert_eq!(top.leaf_policies, vec![RebasePolicy::CheckAncient]); + + // Swap the modes: the INT64 leaf inside `s` is now the only one that needs handling. + let policies = resolve_file_rebase_policies( + &schema, + modes(RebaseReadMode::Exception, RebaseReadMode::Corrected), + ); + let top = wrap_datetime_rebase( + Arc::new(Column::new("ts96", 1)) as Arc, + &schema, + &policies, + ) + .unwrap(); + assert!(top.downcast_ref::().is_some()); + let s = wrap_datetime_rebase( + Arc::new(Column::new("s", 0)) as Arc, + &schema, + &policies, + ) + .unwrap(); + let s = s.downcast_ref::().unwrap(); + assert_eq!( + s.leaf_policies, + vec![ + RebasePolicy::Corrected, + RebasePolicy::Corrected, + RebasePolicy::CheckAncient + ] + ); + } + + #[test] + fn wrap_passes_nested_columns_whose_affected_leaves_are_all_corrected() { + // Date policy Corrected, timestamp policies Legacy, column STRUCT: the + // struct's only rebase-relevant leaf is a date, and the date policy needs no + // handling, so the column must pass through unwrapped instead of being wrapped just + // because SOME policy (timestamps -- absent here) needs handling. + let policies = flat_policies( + RebasePolicy::Corrected, + RebasePolicy::Legacy(WriterTimeZone::Utc), + RebasePolicy::Legacy(WriterTimeZone::Utc), + ); + let schema: SchemaRef = Arc::new(Schema::new(vec![Field::new( + "s", + DataType::Struct(vec![Field::new("d", DataType::Date32, true)].into()), + true, + )])); + let out = wrap_datetime_rebase( + Arc::new(Column::new("s", 0)) as Arc, + &schema, + &policies, + ) + .unwrap(); + assert!(out.downcast_ref::().is_some()); + + // Mirror image: STRUCT under Corrected timestamp policies with a + // non-Corrected date policy passes too. + let policies = flat_policies( + RebasePolicy::CheckAncient, + RebasePolicy::Corrected, + RebasePolicy::Corrected, + ); + let schema: SchemaRef = Arc::new(Schema::new(vec![Field::new( + "s", + DataType::Struct( + vec![Field::new( + "ts", + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + true, + )] + .into(), + ), + true, + )])); + let out = wrap_datetime_rebase( + Arc::new(Column::new("s", 0)) as Arc, + &schema, + &policies, + ) + .unwrap(); + assert!(out.downcast_ref::().is_some()); + } + + #[test] + fn wrap_installs_the_wrapper_on_nested_columns_with_an_affected_leaf() { + // A nested column whose leaves DO include an affected type under a policy that needs + // handling gets the wrapper (not a refusal), across struct, list, and map nesting. + let date_legacy = flat_policies( + RebasePolicy::Legacy(WriterTimeZone::Utc), + RebasePolicy::Corrected, + RebasePolicy::Corrected, + ); + let ts_check = flat_policies( + RebasePolicy::Corrected, + RebasePolicy::CheckAncient, + RebasePolicy::CheckAncient, + ); + let ts_field = Field::new( + "ts", + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + true, + ); + let cases: Vec<(DataType, &FileRebasePolicies, Vec)> = vec![ + ( + DataType::Struct(vec![Field::new("d", DataType::Date32, true)].into()), + &date_legacy, + vec![RebasePolicy::Legacy(WriterTimeZone::Utc)], + ), + ( + DataType::List(Arc::new(Field::new("item", DataType::Date32, true))), + &date_legacy, + vec![RebasePolicy::Legacy(WriterTimeZone::Utc)], + ), + ( + DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct( + vec![ + Field::new("key", DataType::Int64, false), + Field::new("value", DataType::Date32, true), + ] + .into(), + ), + false, + )), + false, + ), + &date_legacy, + vec![ + RebasePolicy::Corrected, + RebasePolicy::Legacy(WriterTimeZone::Utc), + ], + ), + ( + DataType::Struct(vec![ts_field.clone()].into()), + &ts_check, + vec![RebasePolicy::CheckAncient], + ), + ]; + for (dt, policies, expected) in cases { + let schema: SchemaRef = Arc::new(Schema::new(vec![Field::new("n", dt.clone(), true)])); + let wrapped = wrap_datetime_rebase( + Arc::new(Column::new("n", 0)) as Arc, + &schema, + policies, + ) + .unwrap(); + let wrapped = wrapped + .downcast_ref::() + .unwrap_or_else(|| panic!("expected {dt} to be wrapped")); + assert_eq!(wrapped.leaf_policies, expected, "{dt}"); + } + } + + #[test] + fn wrap_passes_nested_columns_with_no_affected_leaves_at_all() { + // A timezone-free timestamp nobody reads as TIMESTAMP and plain types are never + // rebased, so a nested column built only from them passes even when every policy + // needs handling. + let policies = flat_policies( + RebasePolicy::Legacy(WriterTimeZone::Utc), + RebasePolicy::Legacy(WriterTimeZone::Utc), + RebasePolicy::Legacy(WriterTimeZone::Utc), + ); + let schema: SchemaRef = Arc::new(Schema::new(vec![Field::new( + "s", + DataType::Struct( + vec![ + Field::new("i", DataType::Int64, true), + Field::new( + "ntz", + DataType::Timestamp(TimeUnit::Microsecond, None), + true, + ), + ] + .into(), + ), + true, + )])); + let out = wrap_datetime_rebase( + Arc::new(Column::new("s", 0)) as Arc, + &schema, + &policies, + ) + .unwrap(); + assert!(out.downcast_ref::().is_some()); + } + + #[test] + fn nested_struct_list_map_with_modern_and_null_leaves_pass_under_every_policy() { + let date_field = Arc::new(Field::new("d", DataType::Date32, true)); + let ts_field = Arc::new(ts_field(TimeUnit::Microsecond)); + let struct_dt = + DataType::Struct(vec![Arc::clone(&date_field), Arc::clone(&ts_field)].into()); + let struct_arr = StructArray::try_new( + vec![Arc::clone(&date_field), Arc::clone(&ts_field)].into(), + vec![ + Arc::new(Date32Array::from(vec![Some(0), None, Some(19876)])), + ts_array(TimeUnit::Microsecond, vec![Some(-1), None, Some(0)]), + ], + Some(vec![true, true, false].into()), + ) + .unwrap(); + let list_item = Arc::new(Field::new("item", DataType::Date32, true)); + let list_dt = DataType::List(Arc::clone(&list_item)); + let list_arr = ListArray::try_new( + Arc::clone(&list_item), + OffsetBuffer::new(vec![0, 2, 2, 3].into()), + Arc::new(Date32Array::from(vec![Some(0), None, Some(19876)])), + Some(vec![true, false, true].into()), + ) + .unwrap(); + let key_field = Arc::new(Field::new("key", DataType::Int64, false)); + let value_field = Arc::new(Field::new("value", DataType::Date32, true)); + let entries_field = Arc::new(Field::new( + "entries", + DataType::Struct(vec![Arc::clone(&key_field), Arc::clone(&value_field)].into()), + false, + )); + let map_dt = DataType::Map(Arc::clone(&entries_field), false); + let entries = StructArray::try_new( + vec![key_field, value_field].into(), + vec![ + Arc::new(Int64Array::from(vec![1, 2])), + Arc::new(Date32Array::from(vec![Some(19876), None])), + ], + None, + ) + .unwrap(); + let map_arr = MapArray::try_new( + entries_field, + OffsetBuffer::new(vec![0, 1, 2, 2].into()), + entries, + Some(vec![true, true, false].into()), + false, + ) + .unwrap(); + + let cases: Vec<(DataType, ArrayRef)> = vec![ + (struct_dt, Arc::new(struct_arr)), + (list_dt, Arc::new(list_arr)), + (map_dt, Arc::new(map_arr)), + ]; + for policy in [ + RebasePolicy::Legacy(WriterTimeZone::Utc), + RebasePolicy::Legacy(WriterTimeZone::OtherOrUnknown), + RebasePolicy::CheckAncient, + ] { + for (dt, array) in &cases { + let field = Field::new("n", dt.clone(), true); + let expr = rebase_expr(field.clone(), policy); + let out = eval_on(&expr, Arc::clone(array), field) + .unwrap_or_else(|e| panic!("{dt} under {policy:?}: {e}")); + assert_eq!(&out, array, "{dt} under {policy:?} must be the identity"); + } + } + } + + #[test] + fn nested_ancient_date_leaf_rebases_under_legacy_and_errors_under_check_ancient() { + // list>: one ancient leaf among modern and null ones. + let stored = julian_civil_to_day(1500, 1, 1); + let date_field = Arc::new(Field::new("d", DataType::Date32, true)); + let struct_field = Arc::new(Field::new( + "item", + DataType::Struct(vec![Arc::clone(&date_field)].into()), + true, + )); + let dt = DataType::List(Arc::clone(&struct_field)); + let structs = StructArray::try_new( + vec![date_field].into(), + vec![Arc::new(Date32Array::from(vec![ + Some(stored), + None, + Some(19876), + ]))], + Some(vec![true, false, true].into()), + ) + .unwrap(); + let array: ArrayRef = Arc::new( + ListArray::try_new( + Arc::clone(&struct_field), + OffsetBuffer::new(vec![0, 1, 3].into()), + Arc::new(structs), + None, + ) + .unwrap(), + ); + let field = Field::new("n", dt, true); + + let legacy = rebase_expr(field.clone(), RebasePolicy::Legacy(WriterTimeZone::Utc)); + let out = eval_on(&legacy, Arc::clone(&array), field.clone()).unwrap(); + assert_eq!(out.data_type(), field.data_type()); + let out_list = out.as_any().downcast_ref::().unwrap(); + let in_list = array.as_any().downcast_ref::().unwrap(); + assert_eq!(out_list.offsets(), in_list.offsets()); + assert_eq!(out_list.nulls(), in_list.nulls()); + let out_structs = out_list + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(out_structs.nulls(), in_list.values().nulls()); + let dates = out_structs + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(dates.value(0), days_from_civil(1500, 1, 1) as i32); + assert!(dates.is_null(1)); + assert_eq!(dates.value(2), 19876); + + let check = rebase_expr(field.clone(), RebasePolicy::CheckAncient); + let err = eval_on(&check, array, field).unwrap_err().to_string(); + assert!(err.contains("rebase"), "unexpected error: {err}"); + assert!(err.contains("'n'"), "unexpected error: {err}"); + } + + #[test] + fn nested_leaves_each_follow_their_own_policy() { + // struct where the stamp marks the first leaf INT96: + // under datetime CORRECTED + int96 EXCEPTION, an ancient INT64 value passes verbatim + // while an ancient INT96 value in the same struct is refused. + let ts_dt = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())); + let ts96_field = Arc::new(Field::new("ts96", ts_dt.clone(), true)); + let ts_field = Arc::new(Field::new("ts", ts_dt, true)); + let struct_dt = + DataType::Struct(vec![Arc::clone(&ts96_field), Arc::clone(&ts_field)].into()); + let schema: SchemaRef = Arc::new(Schema::new_with_metadata( + vec![Field::new("s", struct_dt.clone(), true)], + spark_metadata(&[(INT96_LEAVES_METADATA_KEY, "2:0")]), + )); + let policies = resolve_file_rebase_policies( + &schema, + modes(RebaseReadMode::Corrected, RebaseReadMode::Exception), + ); + let wrapped = wrap_datetime_rebase( + Arc::new(Column::new("s", 0)) as Arc, + &schema, + &policies, + ) + .unwrap(); + let expr = wrapped.downcast_ref::().unwrap(); + + let ancient = LAST_SWITCH_JULIAN_TS_SECONDS * 1_000_000 - 1; + let build = |ts96: i64, ts: i64| -> ArrayRef { + Arc::new( + StructArray::try_new( + vec![Arc::clone(&ts96_field), Arc::clone(&ts_field)].into(), + vec![ + ts_array(TimeUnit::Microsecond, vec![Some(ts96)]), + ts_array(TimeUnit::Microsecond, vec![Some(ts)]), + ], + None, + ) + .unwrap(), + ) + }; + let field = Field::new("s", struct_dt, true); + let passing = build(0, ancient); + let out = eval_on(expr, Arc::clone(&passing), field.clone()).unwrap(); + assert_eq!(&out, &passing); + let err = eval_on(expr, build(ancient, 0), field) + .unwrap_err() + .to_string(); + assert!(err.contains("rebase"), "unexpected error: {err}"); + } + + /// A physical `s: STRUCT` column holding a modern date and the + /// given timestamp, as (field, type, array). + fn date_ts_struct(ts: i64) -> (FieldRef, DataType, ArrayRef) { + let d_field = Arc::new(Field::new("d", DataType::Date32, true)); + let ts_field = Arc::new(ts_field(TimeUnit::Microsecond)); + let fields: arrow::datatypes::Fields = + vec![Arc::clone(&d_field), Arc::clone(&ts_field)].into(); + let dt = DataType::Struct(fields.clone()); + let array: ArrayRef = Arc::new( + StructArray::try_new( + fields, + vec![ + Arc::new(Date32Array::from(vec![Some(19875)])), + ts_array(TimeUnit::Microsecond, vec![Some(ts)]), + ], + None, + ) + .unwrap(), + ); + (Arc::new(Field::new("s", dt.clone(), true)), dt, array) + } + + fn struct_of(fields: Vec) -> DataType { + DataType::Struct(fields.into()) + } + + #[test] + fn unrequested_struct_leaves_are_never_checked() { + // A metadata-free file with s.d = 2024-06-01 and an ancient s.ts under EXCEPTION read + // modes. Spark's requested schema for `select s.d` is STRUCT, so Spark never decodes + // s.ts; the wrapper sits beneath the schema adapter's struct narrowing and must not + // check the leaf that narrowing drops. + let ancient = LAST_SWITCH_JULIAN_TS_SECONDS * 1_000_000 - 1; + let (field, dt, array) = date_ts_struct(ancient); + let schema = Schema::new(vec![Field::new("s", dt.clone(), true)]); + let policies = resolve_file_rebase_policies(&schema, default_modes()); + assert_eq!(policies.date, RebasePolicy::CheckAncient); + + let requested = struct_of(vec![Field::new("d", DataType::Date32, true)]); + let narrowed = + policies + .clone() + .restrict_to_requested(&schema, &[Some(&requested)], true, false); + assert_eq!(narrowed.unrequested_leaves, vec![1]); + assert!(narrowed.ntz_requested_leaves.is_empty()); + assert!(narrowed.ltz_requested_leaves.is_empty()); + let wrapped = wrap_datetime_rebase( + Arc::new(Column::new("s", 0)) as Arc, + &Arc::new(schema.clone()), + &narrowed, + ) + .unwrap(); + let expr = wrapped.downcast_ref::().unwrap(); + assert_eq!( + expr.leaf_policies, + vec![RebasePolicy::CheckAncient, RebasePolicy::Corrected] + ); + let out = eval_on(expr, Arc::clone(&array), field.as_ref().clone()).unwrap(); + assert_eq!(&out, &array, "the requested modern date passes untouched"); + + // Requesting both leaves (or the whole struct) still refuses the ancient timestamp. + let full = policies + .clone() + .restrict_to_requested(&schema, &[Some(&dt)], true, false); + assert!(full.unrequested_leaves.is_empty()); + for policies in [&policies, &full] { + let wrapped = wrap_datetime_rebase( + Arc::new(Column::new("s", 0)) as Arc, + &Arc::new(schema.clone()), + policies, + ) + .unwrap(); + let expr = wrapped.downcast_ref::().unwrap(); + let err = eval_on(expr, Arc::clone(&array), field.as_ref().clone()) + .unwrap_err() + .to_string(); + assert!(err.contains("rebase"), "unexpected error: {err}"); + } + + // A column with no requested affected leaf at all is not wrapped. + let only_ts_unrequested = policies.clone().restrict_to_requested( + &schema, + &[Some(&struct_of(vec![Field::new( + "d", + DataType::Int32, + true, + )]))], + true, + false, + ); + // (a leaf whose requested type mismatches is still requested -- the cast reads it) + assert!(only_ts_unrequested.unrequested_leaves == vec![1]); + let none_requested = policies.clone().restrict_to_requested( + &schema, + &[Some(&struct_of(vec![Field::new( + "x", + DataType::Int32, + true, + )]))], + true, + false, + ); + assert_eq!(none_requested.unrequested_leaves, vec![0, 1]); + let passthrough = wrap_datetime_rebase( + Arc::new(Column::new("s", 0)) as Arc, + &Arc::new(schema), + &none_requested, + ) + .unwrap(); + assert!(passthrough.downcast_ref::().is_some()); + } + + #[test] + fn unrequested_leaves_inside_lists_are_never_checked() { + // LIST> read as LIST>: the list pairs positionally with the + // requested list and the struct beneath it narrows by name. + let ancient = LAST_SWITCH_JULIAN_TS_SECONDS * 1_000_000 - 1; + let (_, struct_dt, structs) = date_ts_struct(ancient); + let item = Arc::new(Field::new("item", struct_dt, true)); + let list_dt = DataType::List(Arc::clone(&item)); + let array: ArrayRef = Arc::new( + ListArray::try_new(item, OffsetBuffer::new(vec![0, 1].into()), structs, None).unwrap(), + ); + let field = Field::new("l", list_dt.clone(), true); + let schema = Schema::new(vec![field.clone()]); + let policies = resolve_file_rebase_policies(&schema, default_modes()); + + let requested = DataType::List(Arc::new(Field::new( + "element", + struct_of(vec![Field::new("d", DataType::Date32, true)]), + true, + ))); + let narrowed = + policies + .clone() + .restrict_to_requested(&schema, &[Some(&requested)], true, false); + assert_eq!(narrowed.unrequested_leaves, vec![1]); + let wrapped = wrap_datetime_rebase( + Arc::new(Column::new("l", 0)) as Arc, + &Arc::new(schema.clone()), + &narrowed, + ) + .unwrap(); + let expr = wrapped.downcast_ref::().unwrap(); + let out = eval_on(expr, Arc::clone(&array), field.clone()).unwrap(); + assert_eq!(&out, &array); + + let wrapped = wrap_datetime_rebase( + Arc::new(Column::new("l", 0)) as Arc, + &Arc::new(schema), + &policies, + ) + .unwrap(); + let expr = wrapped.downcast_ref::().unwrap(); + let err = eval_on(expr, array, field).unwrap_err().to_string(); + assert!(err.contains("rebase"), "unexpected error: {err}"); + } + + #[test] + fn requested_leaf_narrowing_keeps_int96_ordinals_physical() { + // struct with the stamp marking physical leaf 0 as INT96, read + // as struct only. The attribution must stay keyed on PHYSICAL ordinals: `ts` is + // physical leaf 1 (INT64) even though it is the requested struct's first leaf, so under + // datetime EXCEPTION + int96 CORRECTED it is CheckAncient, and under the swapped modes + // it is Corrected (and the unrequested INT96 leaf is never checked either way). + let ts_dt = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())); + let physical = struct_of(vec![ + Field::new("ts96", ts_dt.clone(), true), + Field::new("ts", ts_dt.clone(), true), + ]); + let schema = Schema::new_with_metadata( + vec![Field::new("s", physical, true)], + spark_metadata(&[(INT96_LEAVES_METADATA_KEY, "2:0")]), + ); + let requested = struct_of(vec![Field::new("ts", ts_dt, true)]); + let leaf_policies_under = |datetime, int96| { + let policies = resolve_file_rebase_policies(&schema, modes(datetime, int96)) + .restrict_to_requested(&schema, &[Some(&requested)], true, false); + assert_eq!(policies.unrequested_leaves, vec![0]); + wrap_datetime_rebase( + Arc::new(Column::new("s", 0)) as Arc, + &Arc::new(schema.clone()), + &policies, + ) + .unwrap() + .downcast_ref::() + .map(|e| e.leaf_policies.clone()) + }; + assert_eq!( + leaf_policies_under(RebaseReadMode::Exception, RebaseReadMode::Corrected), + Some(vec![RebasePolicy::Corrected, RebasePolicy::CheckAncient]) + ); + assert_eq!( + leaf_policies_under(RebaseReadMode::Corrected, RebaseReadMode::Exception), + None + ); + } + + #[test] + fn requested_leaf_narrowing_matches_children_like_the_struct_convert() { + // The mask pairs struct children the way `parquet_convert_struct_to_struct` selects + // them -- by folded name in case-insensitive mode, by Parquet field id when ids are in + // play -- and keeps a child whenever EITHER rule matches, so it can only ever drop + // leaves the narrowing drops too. Shape mismatches and unpaired columns keep every leaf. + use arrow::datatypes::Field as F; + let id = |field: F, id: &str| { + field.with_metadata(HashMap::from([( + parquet::arrow::PARQUET_FIELD_ID_META_KEY.to_string(), + id.to_string(), + )])) + }; + let physical = struct_of(vec![ + id(F::new("A", DataType::Date32, true), "1"), + id(F::new("b", DataType::Date32, true), "2"), + F::new("c", DataType::Date32, true), + F::new("m", DataType::Date32, true), + ]); + let schema = Schema::new(vec![ + Field::new("s", physical, true), + Field::new("d", DataType::Date32, true), + ]); + let policies = resolve_file_rebase_policies(&schema, default_modes()); + let restrict = |requested: &DataType, case_sensitive: bool, use_field_id: bool| { + policies + .clone() + .restrict_to_requested( + &schema, + &[Some(requested), None], + case_sensitive, + use_field_id, + ) + .unrequested_leaves + }; + + // Case-insensitive name match keeps `A` for a requested `a`; case-sensitive drops it. + let by_name = struct_of(vec![F::new("a", DataType::Date32, true)]); + assert_eq!(restrict(&by_name, false, false), vec![1, 2, 3]); + assert_eq!(restrict(&by_name, true, false), vec![0, 1, 2, 3]); + + // Field id 2 selects `b` even though the requested name (`zzz`) matches nothing; the + // unpaired top-level `d` (leaf 4) is never dropped. + let by_id = struct_of(vec![id(F::new("zzz", DataType::Date32, true), "2")]); + assert_eq!(restrict(&by_id, false, true), vec![0, 2, 3]); + // Without field-id matching the id is ignored and nothing pairs. + assert_eq!(restrict(&by_id, false, false), vec![0, 1, 2, 3]); + // Id AND name both count: requested `c` (id 1) keeps physical `A` (id 1) and `c`. + let both = struct_of(vec![id(F::new("c", DataType::Date32, true), "1")]); + assert_eq!(restrict(&both, false, true), vec![1, 3]); + + // A requested type of another shape keeps every leaf (the cast reads them all). + assert_eq!( + restrict(&DataType::Date32, false, false), + Vec::::new() + ); + + // Only the pairings `parquet_convert_array` narrows recurse. A LargeList, a + // FixedSizeList or a dictionary around the struct is handed to arrow's cast (which + // cannot narrow a struct) or passed through whole, so every leaf must stay requested + // even though a plain List around the same struct narrows. + let (_, ts_struct, _) = date_ts_struct(0); + let narrowed_item = struct_of(vec![F::new("d", DataType::Date32, true)]); + let list_schema = |dt: DataType| Schema::new(vec![Field::new("l", dt, true)]); + let list_restrict = |physical: DataType, requested: DataType| { + let schema = list_schema(physical); + resolve_file_rebase_policies(&schema, default_modes()) + .restrict_to_requested(&schema, &[Some(&requested)], true, false) + .unrequested_leaves + }; + let item = |dt: &DataType| Arc::new(F::new("item", dt.clone(), true)); + assert_eq!( + list_restrict( + DataType::List(item(&ts_struct)), + DataType::List(item(&narrowed_item)) + ), + vec![1] + ); + assert_eq!( + list_restrict( + DataType::LargeList(item(&ts_struct)), + DataType::LargeList(item(&narrowed_item)) + ), + Vec::::new() + ); + assert_eq!( + list_restrict( + DataType::FixedSizeList(item(&ts_struct), 1), + DataType::FixedSizeList(item(&narrowed_item), 1) + ), + Vec::::new() + ); + assert_eq!( + list_restrict( + DataType::Dictionary(Box::new(DataType::Int32), Box::new(ts_struct.clone())), + narrowed_item.clone() + ), + Vec::::new() + ); + + // Map: entries pair positionally (key with key, value with value), and a struct value + // narrows by name beneath it -- but only for the same key ordering, the gate + // `parquet_convert_array` puts on its map convert; otherwise every leaf stays. + let entries = |value: DataType, sorted: bool| { + DataType::Map( + Arc::new(F::new( + "entries", + struct_of(vec![ + F::new("key", DataType::Int64, false), + F::new("value", value, true), + ]), + false, + )), + sorted, + ) + }; + let map_schema = Schema::new(vec![Field::new( + "m", + entries(ts_struct.clone(), false), + true, + )]); + let map_policies = resolve_file_rebase_policies(&map_schema, default_modes()); + let requested_value = struct_of(vec![F::new( + "ts", + ts_field(TimeUnit::Microsecond).data_type().clone(), + true, + )]); + let map_restrict = |requested: &DataType| { + map_policies + .clone() + .restrict_to_requested(&map_schema, &[Some(requested)], true, false) + .unrequested_leaves + }; + assert_eq!( + map_restrict(&entries(requested_value.clone(), false)), + vec![1] + ); + assert_eq!( + map_restrict(&entries(requested_value, true)), + Vec::::new() + ); + } + + /// The expression `wrap_datetime_rebase` makes of column 0 of `schema` when the query + /// reads it as `requested` (`None`: the column has no logical counterpart). + fn wrapper_reading( + schema: &Schema, + requested: Option<&DataType>, + session_modes: SessionRebaseModes, + ) -> Arc { + let policies = resolve_file_rebase_policies(schema, session_modes).restrict_to_requested( + schema, + &[requested], + true, + false, + ); + wrap_datetime_rebase( + Arc::new(Column::new(schema.field(0).name(), 0)) as Arc, + &Arc::new(schema.clone()), + &policies, + ) + .unwrap() + } + + /// The leaf policies the wrapper installs on the single column of `schema` when the query + /// reads it as `requested`, or `None` when the column passes through unwrapped. + fn wrapped_policies_reading( + schema: &Schema, + requested: Option<&DataType>, + session_modes: SessionRebaseModes, + ) -> Option> { + wrapper_reading(schema, requested, session_modes) + .downcast_ref::() + .map(|e| e.leaf_policies.clone()) + } + + #[test] + fn ntz_requests_suppress_timestamp_rebase_at_every_depth() { + // Spark's ParquetVectorUpdaterFactory keys on the REQUESTED type: a column read as + // TIMESTAMP_NTZ takes BinaryToSQLTimestampUpdater (INT96) or LongUpdater (INT64), + // neither of which rebases, whatever the read modes say. Both leaves here are + // physically timezone-carrying (leaf 0 INT96 by the stamp, leaf 1 adjusted INT64) + // under EXCEPTION/EXCEPTION, so the physical rule alone would check both. + let ltz = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())); + let ntz = DataType::Timestamp(TimeUnit::Microsecond, None); + let pair = |ts96: &DataType, ts64: &DataType| { + struct_of(vec![ + Field::new("ts96", ts96.clone(), true), + Field::new("ts64", ts64.clone(), true), + ]) + }; + fn list(dt: DataType) -> DataType { + DataType::List(Arc::new(Field::new("item", dt, true))) + } + fn map(dt: DataType) -> DataType { + DataType::Map( + Arc::new(Field::new( + "entries", + struct_of(vec![ + Field::new("key", DataType::Int64, false), + Field::new("value", dt, true), + ]), + false, + )), + false, + ) + } + type Shape = fn(DataType) -> DataType; + let shapes: Vec<(&str, Shape, &str, Vec)> = vec![ + ("struct", |dt| dt, "2:0", vec![]), + ("list", list, "2:0", vec![]), + ("map", map, "3:1", vec![RebasePolicy::Corrected]), + ]; + for (name, shape, stamp, key_leaves) in shapes { + let schema = Schema::new_with_metadata( + vec![Field::new("c", shape(pair(<z, <z)), true)], + spark_metadata(&[(INT96_LEAVES_METADATA_KEY, stamp)]), + ); + let read_as = |ts96: &DataType, ts64: &DataType| { + wrapped_policies_reading(&schema, Some(&shape(pair(ts96, ts64))), default_modes()) + }; + let expect = |leaves: &[RebasePolicy]| { + Some(key_leaves.iter().chain(leaves).copied().collect::>()) + }; + assert_eq!( + read_as(&ntz, &ntz), + None, + "{name}: NTZ requests never rebase" + ); + assert_eq!( + read_as(<z, &ntz), + expect(&[RebasePolicy::CheckAncient, RebasePolicy::Corrected]), + "{name}" + ); + assert_eq!( + read_as(&ntz, <z), + expect(&[RebasePolicy::Corrected, RebasePolicy::CheckAncient]), + "{name}" + ); + assert_eq!( + read_as(<z, <z), + expect(&[RebasePolicy::CheckAncient, RebasePolicy::CheckAncient]), + "{name}" + ); + // An unpaired column keeps the physical rule. + assert_eq!( + wrapped_policies_reading(&schema, None, default_modes()), + expect(&[RebasePolicy::CheckAncient, RebasePolicy::CheckAncient]), + "{name}" + ); + } + + // End to end on the struct: an ancient adjusted INT64 value passes when its leaf is + // read as TIMESTAMP_NTZ and is refused when it is read as TIMESTAMP. + let schema = Schema::new_with_metadata( + vec![Field::new("c", pair(<z, <z), true)], + spark_metadata(&[(INT96_LEAVES_METADATA_KEY, "2:0")]), + ); + let ancient = LAST_SWITCH_JULIAN_TS_SECONDS * 1_000_000 - 1; + let array: ArrayRef = Arc::new( + StructArray::try_new( + vec![ + Arc::new(Field::new("ts96", ltz.clone(), true)), + Arc::new(Field::new("ts64", ltz.clone(), true)), + ] + .into(), + vec![ + ts_array(TimeUnit::Microsecond, vec![Some(0)]), + ts_array(TimeUnit::Microsecond, vec![Some(ancient)]), + ], + None, + ) + .unwrap(), + ); + let field = schema.field(0).as_ref().clone(); + let wrapped_reading = |requested: DataType| { + let policies = resolve_file_rebase_policies(&schema, default_modes()) + .restrict_to_requested(&schema, &[Some(&requested)], true, false); + assert!(policies.unrequested_leaves.is_empty()); + assert!(policies.ltz_requested_leaves.is_empty()); + wrap_datetime_rebase( + Arc::new(Column::new("c", 0)) as Arc, + &Arc::new(schema.clone()), + &policies, + ) + .unwrap() + }; + let mixed = wrapped_reading(pair(<z, &ntz)); + let expr = mixed.downcast_ref::().unwrap(); + assert_eq!( + expr.leaf_policies, + vec![RebasePolicy::CheckAncient, RebasePolicy::Corrected] + ); + let out = eval_on(expr, Arc::clone(&array), field.clone()).unwrap(); + assert_eq!(&out, &array); + let both = wrapped_reading(pair(<z, <z)); + let expr = both.downcast_ref::().unwrap(); + let err = eval_on(expr, array, field).unwrap_err().to_string(); + assert!(err.contains("rebase"), "unexpected error: {err}"); + } + + #[test] + fn ltz_requests_on_tz_free_leaves_follow_the_datetime_spec() { + // A physical INT64 timestamp with isAdjustedToUTC=false surfaces as a timezone-free + // arrow timestamp. Spark's INT64 branch checks only the unit (isTimestampTypeMatched), + // so reading it as TIMESTAMP takes LongWithRebaseUpdater under the datetime spec, and + // reading it as TIMESTAMP_NTZ takes LongUpdater, which never rebases. + let ltz = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())); + for unit in [TimeUnit::Microsecond, TimeUnit::Millisecond] { + let ntz = DataType::Timestamp(unit, None); + let schema = Schema::new(vec![Field::new("ts", ntz.clone(), true)]); + let read_as = |requested: Option<&DataType>, datetime, int96| { + wrapped_policies_reading(&schema, requested, modes(datetime, int96)) + }; + assert_eq!( + read_as( + Some(<z), + RebaseReadMode::Exception, + RebaseReadMode::Corrected + ), + Some(vec![RebasePolicy::CheckAncient]), + "{unit:?}" + ); + // The datetime spec, not the INT96 one. + assert_eq!( + read_as( + Some(<z), + RebaseReadMode::Corrected, + RebaseReadMode::Exception + ), + None, + "{unit:?}" + ); + assert_eq!( + read_as( + Some(&ntz), + RebaseReadMode::Exception, + RebaseReadMode::Exception + ), + None, + "{unit:?}" + ); + // An unpaired column keeps the physical rule: nothing reads it as TIMESTAMP. + assert_eq!( + read_as(None, RebaseReadMode::Exception, RebaseReadMode::Exception), + None, + "{unit:?}" + ); + } + + // Nested: the leaf is attributed by its physical ordinal beneath the struct. + let ntz = DataType::Timestamp(TimeUnit::Microsecond, None); + let nested = Schema::new(vec![Field::new( + "s", + struct_of(vec![ + Field::new("i", DataType::Int64, true), + Field::new("ts", ntz.clone(), true), + ]), + true, + )]); + assert_eq!( + wrapped_policies_reading( + &nested, + Some(&struct_of(vec![ + Field::new("i", DataType::Int64, true), + Field::new("ts", ltz.clone(), true), + ])), + modes(RebaseReadMode::Exception, RebaseReadMode::Corrected), + ), + Some(vec![RebasePolicy::Corrected, RebasePolicy::CheckAncient]) + ); + + // A stamp naming the leaf INT96 sends it to the INT96 spec, like any INT96 leaf. + let stamped = Schema::new_with_metadata( + vec![Field::new("ts", ntz.clone(), true)], + spark_metadata(&[(INT96_LEAVES_METADATA_KEY, "1:0")]), + ); + assert_eq!( + wrapped_policies_reading( + &stamped, + Some(<z), + modes(RebaseReadMode::Corrected, RebaseReadMode::Exception), + ), + Some(vec![RebasePolicy::CheckAncient]) + ); + + // LEGACY with a UTC writer zone rebases the value; the output stays timezone-free (the + // wrapper sits beneath the adapter's cast to the requested type). + const MICROS_PER_DAY: i64 = 86_400_000_000; + let field = Field::new("ts", ntz.clone(), true); + let legacy = Schema::new_with_metadata( + vec![field.clone()], + spark_metadata(&[(SPARK_TIMEZONE_KEY, "UTC")]), + ); + let policies = resolve_file_rebase_policies( + &legacy, + modes(RebaseReadMode::Legacy, RebaseReadMode::Exception), + ) + .restrict_to_requested(&legacy, &[Some(<z)], true, false); + assert_eq!(policies.ltz_requested_leaves, vec![0]); + assert!(policies.ntz_requested_leaves.is_empty()); + let wrapped = wrap_datetime_rebase( + Arc::new(Column::new("ts", 0)) as Arc, + &Arc::new(legacy.clone()), + &policies, + ) + .unwrap(); + let expr = wrapped.downcast_ref::().unwrap(); + assert_eq!( + expr.leaf_policies, + vec![RebasePolicy::Legacy(WriterTimeZone::Utc)] + ); + let time_of_day = (12i64 * 3600 + 34 * 60 + 56) * 1_000_000; + let stored = julian_civil_to_day(1500, 1, 1) as i64 * MICROS_PER_DAY + time_of_day; + let input: ArrayRef = Arc::new(TimestampMicrosecondArray::from(vec![ + Some(stored), + None, + Some(1_700_000_000_000_000), + ])); + let out = eval_on(expr, input, field).unwrap(); + assert_eq!(out.data_type(), &ntz); + let out = out + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + out.value(0), + days_from_civil(1500, 1, 1) * MICROS_PER_DAY + time_of_day + ); + assert!(out.is_null(1)); + assert_eq!(out.value(2), 1_700_000_000_000_000); + } + + #[test] + fn date_leaves_requested_as_ntz_keep_the_date_policy() { + // Spark 4.x reads DATE as TIMESTAMP_NTZ through DateToTimestampNTZWithRebaseUpdater, + // under the datetime spec exactly like DATE itself (3.x has no such arm), so the + // requested type changes nothing for a Date32 leaf. + let ntz = DataType::Timestamp(TimeUnit::Microsecond, None); + let plain = Schema::new(vec![Field::new("d", DataType::Date32, true)]); + let legacy = Schema::new_with_metadata( + vec![Field::new("d", DataType::Date32, true)], + spark_metadata(&[ + (SPARK_VERSION_METADATA_KEY, "3.5.9"), + (SPARK_LEGACY_DATETIME_KEY, ""), + (SPARK_TIMEZONE_KEY, "UTC"), + ]), + ); + for requested in [DataType::Date32, ntz] { + assert_eq!( + wrapped_policies_reading(&plain, Some(&requested), default_modes()), + Some(vec![RebasePolicy::CheckAncient]), + "{requested}" + ); + assert_eq!( + wrapped_policies_reading(&legacy, Some(&requested), default_modes()), + Some(vec![RebasePolicy::Legacy(WriterTimeZone::Utc)]), + "{requested}" + ); + } + } + + #[test] + fn numeric_requests_on_temporal_leaves_never_rebase() { + // Spark's ParquetVectorUpdaterFactory picks the numeric updaters (IntegerUpdater, + // ByteUpdater, ShortUpdater for INT32, plus IntegerToLongUpdater and + // IntegerToDoubleUpdater on 4.x; LongUpdater for INT64) before any DATE or TIMESTAMP + // arm, so a temporal leaf read as a number returns the stored value with no calendar + // check. Each pair is checked under a metadata-free file with EXCEPTION and LEGACY + // session modes, and under a Spark 2.4 file whose footer forces LEGACY. + let utc = || Some("UTC".into()); + let cases: Vec<(DataType, Vec)> = vec![ + ( + DataType::Timestamp(TimeUnit::Nanosecond, utc()), + vec![DataType::Int64], + ), + ( + DataType::Timestamp(TimeUnit::Microsecond, utc()), + vec![DataType::Int64], + ), + ( + DataType::Timestamp(TimeUnit::Millisecond, utc()), + vec![DataType::Int64], + ), + ( + DataType::Timestamp(TimeUnit::Microsecond, None), + vec![DataType::Int64], + ), + ( + DataType::Date32, + vec![ + DataType::Int32, + DataType::Int8, + DataType::Int16, + DataType::Int64, + DataType::Float64, + ], + ), + ]; + let legacy_metadata = spark_metadata(&[ + (SPARK_VERSION_METADATA_KEY, "2.4.8"), + (SPARK_TIMEZONE_KEY, "UTC"), + ]); + let session_modes = [ + modes(RebaseReadMode::Exception, RebaseReadMode::Exception), + modes(RebaseReadMode::Legacy, RebaseReadMode::Legacy), + ]; + for (physical, numeric) in cases { + let plain = Schema::new(vec![Field::new("v", physical.clone(), true)]); + let legacy = Schema::new_with_metadata( + vec![Field::new("v", physical.clone(), true)], + legacy_metadata.clone(), + ); + for requested in &numeric { + for session in session_modes { + for schema in [&plain, &legacy] { + assert_eq!( + wrapped_policies_reading(schema, Some(requested), session), + None, + "{physical} read as {requested} under {session:?}" + ); + } + } + } + // Control: the same leaf read as its own temporal type keeps its policy. A + // timezone-free leaf is read as TIMESTAMP, the request that rebases it. + let temporal = match &physical { + DataType::Timestamp(unit, None) => DataType::Timestamp(*unit, utc()), + other => other.clone(), + }; + assert_eq!( + wrapped_policies_reading(&plain, Some(&temporal), session_modes[0]), + Some(vec![RebasePolicy::CheckAncient]), + "{physical} read as {temporal}" + ); + assert_eq!( + wrapped_policies_reading(&legacy, Some(&temporal), session_modes[0]), + Some(vec![RebasePolicy::Legacy(WriterTimeZone::Utc)]), + "{physical} read as {temporal}" + ); + } + + // No Spark updater reads a timestamp as anything narrower than BIGINT, so those + // requests are not numeric reads and keep the physical rule. + let ts = Schema::new(vec![Field::new( + "v", + DataType::Timestamp(TimeUnit::Microsecond, utc()), + true, + )]); + for requested in [ + DataType::Int8, + DataType::Int16, + DataType::Int32, + DataType::Float64, + ] { + let policies = resolve_file_rebase_policies(&ts, default_modes()) + .restrict_to_requested(&ts, &[Some(&requested)], true, false); + assert!(policies.numeric_requested_leaves.is_empty(), "{requested}"); + assert_eq!( + wrapped_policies_reading(&ts, Some(&requested), default_modes()), + Some(vec![RebasePolicy::CheckAncient]), + "{requested}" + ); + } + } + + #[test] + fn numeric_requests_suppress_rebase_inside_nested_columns() { + // The pairing recurses through structs, lists and maps, so a nested temporal leaf read + // as a number is the identity while its temporal siblings keep their policy. + let ts = DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())); + let pair = |d: DataType, t: DataType| { + struct_of(vec![Field::new("d", d, true), Field::new("t", t, true)]) + }; + let list = |dt: DataType| DataType::List(Arc::new(Field::new("item", dt, true))); + let map = |dt: DataType| { + DataType::Map( + Arc::new(Field::new( + "entries", + struct_of(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", dt, true), + ]), + false, + )), + false, + ) + }; + let physical = struct_of(vec![ + Field::new("s", pair(DataType::Date32, ts.clone()), true), + Field::new("l", list(ts.clone()), true), + Field::new("m", map(DataType::Date32), true), + ]); + let requested = |s: DataType, l: DataType, m: DataType| { + struct_of(vec![ + Field::new("s", s, true), + Field::new("l", list(l), true), + Field::new("m", map(m), true), + ]) + }; + let schema = Schema::new(vec![Field::new("c", physical.clone(), true)]); + let check = RebasePolicy::CheckAncient; + let pass = RebasePolicy::Corrected; + assert_eq!( + wrapped_policies_reading( + &schema, + Some(&requested( + pair(DataType::Int32, DataType::Int64), + DataType::Int64, + DataType::Int32, + )), + default_modes(), + ), + None + ); + assert_eq!( + wrapped_policies_reading( + &schema, + Some(&requested( + pair(DataType::Date32, DataType::Int64), + ts.clone(), + DataType::Int32, + )), + default_modes(), + ), + Some(vec![check, pass, check, pass, pass]) + ); + assert_eq!( + wrapped_policies_reading( + &schema, + Some(&requested( + pair(DataType::Int32, ts.clone()), + DataType::Int64, + DataType::Date32, + )), + default_modes(), + ), + Some(vec![pass, check, pass, pass, check]) + ); + // Control: read as stored, every temporal leaf keeps its policy. + assert_eq!( + wrapped_policies_reading(&schema, Some(&physical), default_modes()), + Some(vec![check, check, check, pass, check]) + ); + } + + #[test] + fn int32_leaves_read_as_date_take_the_date_policy() { + // Spark's ParquetVectorUpdaterFactory picks the INT32 updater from the requested type + // alone: a DateType request gets IntegerWithRebaseUpdater under the datetime spec + // (IntegerUpdater under CORRECTED) whatever the column's annotation, so a plain INT32 + // column read as DATE follows the file's date policy exactly like a DATE column. + let plain = Schema::new(vec![Field::new("v", DataType::Int32, true)]); + let legacy_footer = Schema::new_with_metadata( + vec![Field::new("v", DataType::Int32, true)], + spark_metadata(&[ + (SPARK_VERSION_METADATA_KEY, "2.4.8"), + (SPARK_TIMEZONE_KEY, "UTC"), + ]), + ); + let date = Some(&DataType::Date32); + let corrected = modes(RebaseReadMode::Corrected, RebaseReadMode::Corrected); + let legacy = modes(RebaseReadMode::Legacy, RebaseReadMode::Legacy); + let exception = modes(RebaseReadMode::Exception, RebaseReadMode::Exception); + assert_eq!( + wrapped_policies_reading(&plain, date, legacy), + Some(vec![RebasePolicy::Legacy(WriterTimeZone::OtherOrUnknown)]) + ); + assert_eq!( + wrapped_policies_reading(&plain, date, exception), + Some(vec![RebasePolicy::CheckAncient]) + ); + assert_eq!(wrapped_policies_reading(&plain, date, corrected), None); + for session in [corrected, legacy, exception] { + // A legacy Spark writer's footer decides over every session mode. + assert_eq!( + wrapped_policies_reading(&legacy_footer, date, session), + Some(vec![RebasePolicy::Legacy(WriterTimeZone::Utc)]), + "{session:?}" + ); + // Control: read as INT (IntegerUpdater), as BIGINT (refused for DATE by Spark's + // INT64 arm, and never a date request), or with no logical counterpart, the stored + // value is never rebased. + for schema in [&plain, &legacy_footer] { + for requested in [Some(&DataType::Int32), Some(&DataType::Int64), None] { + assert_eq!( + wrapped_policies_reading(schema, requested, session), + None, + "{requested:?} under {session:?}" + ); + } + } + } + // INT64 storage has no DATE updater in Spark, so a date request keeps the identity. + let int64 = Schema::new(vec![Field::new("v", DataType::Int64, true)]); + assert_eq!(wrapped_policies_reading(&int64, date, exception), None); + // Nor is an INT_8 leaf, whose Int8 values the scan cannot cast to Date32. + let int8 = Schema::new(vec![Field::new("v", DataType::Int8, true)]); + assert!(resolve_file_rebase_policies(&int8, exception) + .restrict_to_requested(&int8, &[date], true, false) + .date_requested_leaves + .is_empty()); + } + + #[test] + fn int32_leaves_read_as_date_rebase_the_stored_days_before_the_cast() { + // The wrapper sits beneath the schema adapter's INT32 -> DATE cast, so it sees and + // returns the raw Int32 days. + let schema = Schema::new(vec![Field::new("v", DataType::Int32, true)]); + let stored = julian_civil_to_day(1500, 1, 1); + assert_eq!(stored, -171655); + let batch = RecordBatch::try_new( + Arc::new(schema.clone()), + vec![Arc::new(Int32Array::from(vec![Some(stored), None, Some(19876)])) as ArrayRef], + ) + .unwrap(); + + let legacy = wrapper_reading( + &schema, + Some(&DataType::Date32), + modes(RebaseReadMode::Legacy, RebaseReadMode::Legacy), + ); + let rebased = legacy.evaluate(&batch).unwrap().into_array(3).unwrap(); + let rebased = rebased.as_primitive::(); + assert_eq!(rebased.value(0), days_from_civil(1500, 1, 1) as i32); + assert!(rebased.is_null(1)); + assert_eq!(rebased.value(2), 19876); + + let exception = wrapper_reading( + &schema, + Some(&DataType::Date32), + modes(RebaseReadMode::Exception, RebaseReadMode::Exception), + ); + let err = exception.evaluate(&batch).unwrap_err().to_string(); + assert!( + err.contains("Native scan cannot rebase") && err.contains("'v'"), + "unexpected error: {err}" + ); + let modern = RecordBatch::try_new( + Arc::new(schema.clone()), + vec![Arc::new(Int32Array::from(vec![Some(19876), None])) as ArrayRef], + ) + .unwrap(); + assert!(exception.evaluate(&modern).is_ok()); + + let corrected = wrapper_reading( + &schema, + Some(&DataType::Date32), + modes(RebaseReadMode::Corrected, RebaseReadMode::Corrected), + ); + assert!(corrected + .downcast_ref::() + .is_none()); + } + + #[test] + fn int32_leaves_read_as_date_inside_nested_columns_take_the_date_policy() { + // The pairing recurses through structs and lists, so a nested INT32 leaf read as DATE + // takes the date policy while an INT32 sibling read as INT stays the identity. + let list = |dt: DataType| DataType::List(Arc::new(Field::new("item", dt, true))); + let physical = struct_of(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + Field::new("l", list(DataType::Int32), true), + ]); + let requested = struct_of(vec![ + Field::new("a", DataType::Date32, true), + Field::new("b", DataType::Int32, true), + Field::new("l", list(DataType::Date32), true), + ]); + let schema = Schema::new(vec![Field::new("c", physical.clone(), true)]); + let legacy = modes(RebaseReadMode::Legacy, RebaseReadMode::Legacy); + let exception = modes(RebaseReadMode::Exception, RebaseReadMode::Exception); + let check = RebasePolicy::CheckAncient; + let rebase = RebasePolicy::Legacy(WriterTimeZone::OtherOrUnknown); + let pass = RebasePolicy::Corrected; + assert_eq!( + wrapped_policies_reading(&schema, Some(&requested), exception), + Some(vec![check, pass, check]) + ); + assert_eq!( + wrapped_policies_reading(&schema, Some(&requested), legacy), + Some(vec![rebase, pass, rebase]) + ); + assert_eq!( + wrapped_policies_reading( + &schema, + Some(&requested), + modes(RebaseReadMode::Corrected, RebaseReadMode::Corrected) + ), + None + ); + // Control: read as stored, no leaf is a date. + assert_eq!( + wrapped_policies_reading(&schema, Some(&physical), exception), + None + ); + + let stored = julian_civil_to_day(1500, 1, 1); + let DataType::Struct(fields) = &physical else { + unreachable!() + }; + let column = StructArray::try_new( + fields.clone(), + vec![ + Arc::new(Int32Array::from(vec![Some(stored)])) as ArrayRef, + Arc::new(Int32Array::from(vec![Some(stored)])) as ArrayRef, + Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![Some(stored), None]), + ])) as ArrayRef, + ], + None, + ) + .unwrap(); + let batch = + RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(column) as ArrayRef]) + .unwrap(); + let rebased = wrapper_reading(&schema, Some(&requested), legacy) + .evaluate(&batch) + .unwrap() + .into_array(1) + .unwrap(); + let rebased = rebased.as_struct(); + let proleptic = days_from_civil(1500, 1, 1) as i32; + assert_eq!( + rebased.column(0).as_primitive::().value(0), + proleptic + ); + assert_eq!( + rebased.column(1).as_primitive::().value(0), + stored + ); + let items = rebased + .column(2) + .as_list::() + .values() + .as_primitive::(); + assert_eq!(items.value(0), proleptic); + assert!(items.is_null(1)); + } +} 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 dbde337306e..55263dbaff8 100644 --- a/native/core/src/parquet/eager_page_index_reader_factory.rs +++ b/native/core/src/parquet/eager_page_index_reader_factory.rs @@ -58,11 +58,26 @@ //! 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. //! +//! The third is the INT96 leaf stamp, `with_int96_leaf_stamp`, enabled by rebase-aware scans: the +//! factory stamps each unencrypted file's INT96 leaf ordinals into the in-memory copy of its footer +//! key-value metadata -- `datetime_rebase::stamp_int96_leaves`, derived from the footer's own +//! `SchemaDescriptor` -- and caches the stamped copy in place of the plain one. parquet-rs copies +//! every key-value pair into the arrow schema it derives from the metadata, which is the only +//! per-file channel DataFusion's opener gives the expression adapter; the stamp is how the adapter +//! tells INT96 timestamp columns from INT64 ones after both were coerced to the same arrow type. +//! The rebuild happens once per file per cache lifetime (later opens find the stamp already +//! present); encrypted opens are left untouched because the parquet API cannot carry a file +//! decryptor across the rebuild. `FileMetadataCache` is keyed by object path and shared by every +//! scan of one `RuntimeEnv`, so a plain (non-stamping) scan of the same file in the same plan sees +//! the stamped copy too; nothing outside the rebase path reads the key, and the copy is otherwise +//! identical. +//! //! Conversion checks also use the raw footer to preserve legacy LIST clipping failures. With //! row filters enabled, a deferred conversion failure is raised on the first data-page request. //! Footer, Bloom-filter and page-index reads still proceed, so format pruning can discard a file //! without decoding it, but row selection cannot hide the failure. +use crate::parquet::datetime_rebase::stamp_int96_leaves; use crate::parquet::parquet_support::SparkParquetOptions; use crate::parquet::schema_adapter::{ check_file_conversions, RejectOnNonEmpty, REPEATED_PRIMITIVE_KEY, @@ -71,11 +86,13 @@ use arrow::datatypes::{DataType, FieldRef, Schema, SchemaRef}; use async_trait::async_trait; use bytes::Bytes; use datafusion::common::{DataFusionError, Result as DFResult}; -use datafusion::datasource::physical_plan::parquet::metadata::DFParquetMetadata; +use datafusion::datasource::physical_plan::parquet::metadata::{ + CachedParquetMetaData, DFParquetMetadata, +}; use datafusion::datasource::physical_plan::parquet::{ ParquetFileMetrics, ParquetFileReaderFactory, }; -use datafusion::execution::cache::cache_manager::FileMetadataCache; +use datafusion::execution::cache::cache_manager::{CachedFileMetadataEntry, FileMetadataCache}; use datafusion::physical_plan::metrics::{ Count, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, MetricType, }; @@ -189,6 +206,7 @@ pub struct EagerPageIndexReaderFactory { // Refuse a file whose Parquet schema carries no field id, as Spark's `ParquetReadSupport` // does when the requested schema carries one and `ignoreMissing` is not set. require_field_ids: bool, + stamp_int96_leaves: bool, conversion_check: Option<(SchemaRef, SparkParquetOptions, bool)>, deferred_rejections: Arc>>, } @@ -220,6 +238,7 @@ impl EagerPageIndexReaderFactory { scan_io_metrics, spark_variant_schema: false, require_field_ids: false, + stamp_int96_leaves: false, conversion_check: None, deferred_rejections: Arc::new(Mutex::new(HashMap::new())), } @@ -248,6 +267,13 @@ impl EagerPageIndexReaderFactory { self.require_field_ids = enabled; self } + + /// Whether readers stamp each unencrypted file's INT96 leaf ordinals into its metadata + /// (see the module docs). Off by default. + pub fn with_int96_leaf_stamp(mut self, enabled: bool) -> Self { + self.stamp_int96_leaves = enabled; + self + } } /// True when `node` or any node under it carries a field id, the way Spark's @@ -296,6 +322,7 @@ impl ParquetFileReaderFactory for EagerPageIndexReaderFactory { metadata_size_hint, spark_variant_schema: self.spark_variant_schema, require_field_ids: self.require_field_ids, + stamp_int96_leaves: self.stamp_int96_leaves, conversion_check: self.conversion_check.clone(), deferred_rejections: Arc::clone(&self.deferred_rejections), deferred_rejection, @@ -315,6 +342,7 @@ struct EagerPageIndexReader { metadata_size_hint: Option, spark_variant_schema: bool, require_field_ids: bool, + stamp_int96_leaves: bool, conversion_check: Option<(SchemaRef, SparkParquetOptions, bool)>, deferred_rejections: Arc>>, deferred_rejection: Option, @@ -571,6 +599,7 @@ impl AsyncFileReader for EagerPageIndexReader { let scan_io_metrics = Arc::clone(&self.scan_io_metrics); let spark_variant_schema = self.spark_variant_schema; let require_field_ids = self.require_field_ids; + let stamp_enabled = self.stamp_int96_leaves; async move { let file_decryption_properties = options .and_then(|o| o.file_decryption_properties()) @@ -603,7 +632,7 @@ impl AsyncFileReader for EagerPageIndexReader { let metadata = DFParquetMetadata::new(&metadata_store, &object_meta) .with_decryption_properties(file_decryption_properties) - .with_file_metadata_cache(Some(metadata_cache)) + .with_file_metadata_cache(Some(Arc::clone(&metadata_cache))) .with_metadata_size_hint(metadata_size_hint) .with_page_index_policy(page_index_policy) .fetch_metadata() @@ -642,6 +671,32 @@ impl AsyncFileReader for EagerPageIndexReader { }, ))); } + // Stamp before the Variant rewrite so the shared cache keeps the footer as read + // plus the stamp, which the rewrite leaves in place; the rewrite is per open and + // never written back. Encrypted opens (`!cache_enabled`) are never stamped: + // nothing is cached for them and the rebuild cannot carry a file decryptor. + let metadata = if stamp_enabled && cache_enabled { + // First open of this file since the cache last held it: rebuild once with the + // stamp and replace the cached plain copy so later opens skip the rebuild. + // Same entry shape `DFParquetMetadata::cache_metadata` stores, so cache + // validation and page-index reuse behave identically. + match stamp_int96_leaves(&metadata) { + None => metadata, + Some(stamped) => { + let stamped = Arc::new(stamped); + metadata_cache.put( + &object_meta.location, + CachedFileMetadataEntry::new( + object_meta.clone(), + Arc::new(CachedParquetMetaData::new(Arc::clone(&stamped))), + ), + ); + stamped + } + } + } else { + metadata + }; let metadata = if spark_variant_schema { with_spark_arrow_schema(metadata)? } else { diff --git a/native/core/src/parquet/mod.rs b/native/core/src/parquet/mod.rs index 7930320d148..9c63b34bff2 100644 --- a/native/core/src/parquet/mod.rs +++ b/native/core/src/parquet/mod.rs @@ -24,5 +24,6 @@ pub mod schema_adapter; pub mod util; mod cast_column; +mod datetime_rebase; mod name_fold; pub(crate) mod objectstore; diff --git a/native/core/src/parquet/objectstore/s3.rs b/native/core/src/parquet/objectstore/s3.rs index 3ef20195354..ba5c74498d2 100644 --- a/native/core/src/parquet/objectstore/s3.rs +++ b/native/core/src/parquet/objectstore/s3.rs @@ -436,6 +436,32 @@ pub(super) fn get_config_trimmed<'a>( get_config(configs, bucket, property).map(|s| s.trim()) } +/// Every `fs.s3a.*` property suffix (without the `fs.s3a.` prefix) this module resolves via +/// [`get_config`]/[`get_config_trimmed`], i.e. every Hadoop S3A config key native's S3 client +/// actually reads. `#[cfg(test)]`-only. +/// +/// SYNC NOTE: a new call site means adding its property here and `fs.s3a.` to +/// `S3ConfigKeyConsumers` in +/// `contrib/delta-spark/src/main/scala/org/apache/comet/contrib/delta/DeltaScanSupport.scala`. +/// `native_s3a_config_properties_matches_call_sites` keeps this list equal to the call sites, +/// `delta_contrib_compares_every_native_s3a_property` checks the Scala block covers it, and +/// `DeltaScanContribSuite` checks it against the compiled `AllS3ConfigKeys`. +#[cfg(test)] +pub(super) const NATIVE_S3A_CONFIG_PROPERTIES: &[&str] = &[ + "endpoint.region", + "path.style.access", + "endpoint", + "requester.pays.enabled", + "comet.credential.provider.class", + "aws.credentials.provider", + "access.key", + "secret.key", + "session.token", + "assumed.role.credentials.provider", + "assumed.role.arn", + "assumed.role.session.name", +]; + /// Activation key (without `fs.s3a.` prefix) naming the vendor `CometS3CredentialProvider` FQCN. /// Per-bucket override is honored via [`get_config_trimmed`]. const PROVIDER_CLASS_PROPERTY: &str = "comet.credential.provider.class"; @@ -1002,10 +1028,161 @@ impl CredentialProviderMetadata { #[cfg(test)] mod tests { + use std::collections::BTreeSet; use std::sync::atomic::{AtomicI32, Ordering}; use super::*; + /// Discovery-harness test (see `NATIVE_S3A_CONFIG_PROPERTIES`'s doc): mechanically re-derives + /// the set of `fs.s3a.*` property suffixes this file actually resolves by scanning this + /// file's OWN source text (via `include_str!`) for every `get_config(configs, bucket, ...)`/ + /// `get_config_trimmed(configs, bucket, ...)` call site, resolving an identifier argument + /// (e.g. `PROVIDER_CLASS_PROPERTY`) through its own `const NAME: &str = "..."` definition, and + /// asserts the result is EXACTLY `NATIVE_S3A_CONFIG_PROPERTIES`. This fails loudly the moment + /// a call site is added, removed, or its literal changes without updating that constant. + /// It guards against a config key read on one side of the Scala/Rust boundary but not + /// compared on the other, which lets Hadoop and native resolve that key differently. + /// + /// The `configs, property` call inside `get_config_trimmed`'s own body (a passthrough of its + /// own `property` parameter, not a call site naming a fixed config key) is deliberately + /// excluded by name. + #[test] + fn native_s3a_config_properties_matches_call_sites() { + let full_source = include_str!("s3.rs"); + // Scan only the non-test portion of this file: the test module below (this very test) + // necessarily contains the pattern strings themselves as strings, which would otherwise + // make the scan match itself and capture garbage. + let test_mod_start = full_source + .find("#[cfg(test)]\nmod tests {") + .expect("this file must contain a `#[cfg(test)] mod tests {` block"); + let source = &full_source[..test_mod_start]; + let mut found: BTreeSet = BTreeSet::new(); + + for pattern in [ + "get_config_trimmed(configs, bucket, ", + "get_config(configs, bucket, ", + ] { + let mut search_start = 0usize; + while let Some(rel_idx) = source[search_start..].find(pattern) { + let start = search_start + rel_idx + pattern.len(); + let end = start + + source[start..] + .find(')') + .expect("unterminated get_config(_trimmed) call in source scan"); + let arg = source[start..end].trim(); + search_start = end + 1; + + if arg == "property" { + // get_config_trimmed's own passthrough of its `property` parameter -- not a + // call site naming a fixed config key. + continue; + } + + let literal = if let Some(stripped) = arg.strip_prefix('"') { + stripped + .strip_suffix('"') + .unwrap_or_else(|| panic!("malformed string literal argument: {arg}")) + .to_string() + } else { + // Identifier argument (e.g. PROVIDER_CLASS_PROPERTY): resolve via its own + // `const NAME: &str = "value";` definition elsewhere in this file. + let const_decl = format!("const {arg}: &str = \""); + let decl_start = source.find(&const_decl).unwrap_or_else(|| { + panic!( + "no `const {arg}: &str = \"...\";` definition found for identifier \ + argument passed to get_config/get_config_trimmed -- update this \ + test's resolution logic or the source" + ) + }) + const_decl.len(); + let decl_end = source[decl_start..] + .find('"') + .expect("unterminated const string literal") + + decl_start; + source[decl_start..decl_end].to_string() + }; + found.insert(literal); + } + } + + let expected: BTreeSet = NATIVE_S3A_CONFIG_PROPERTIES + .iter() + .map(|s| s.to_string()) + .collect(); + + assert_eq!( + found, expected, + "NATIVE_S3A_CONFIG_PROPERTIES must exactly match every property name passed to \ + get_config/get_config_trimmed in this file -- update the constant (and keep \ + DeltaScanSupport.scala's AllS3ConfigKeys in sync, see that constant's SYNC NOTE) \ + when a call site changes" + ); + } + + /// Checks that the Delta contrib's `S3ConfigKeyConsumers` (`DeltaScanSupport.scala`) compares + /// every property in `NATIVE_S3A_CONFIG_PROPERTIES`, on every PR that touches native code + /// (the contrib suites do not). Parses the Scala block as text: each `"fs.s3a."` + /// literal followed by `->` up to the first `)`. `DeltaScanContribSuite` checks the same + /// parse against the compiled `AllS3ConfigKeys`. + #[test] + fn delta_contrib_compares_every_native_s3a_property() { + let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join( + "../../contrib/delta-spark/src/main/scala/org/apache/comet/contrib/delta/\ + DeltaScanSupport.scala", + ); + let scala = std::fs::read_to_string(&path) + .unwrap_or_else(|e| panic!("cannot read {}: {e}", path.display())); + + let marker = "S3ConfigKeyConsumers: Seq[(String, S3ConfigConsumer)] = Seq("; + let start = scala.find(marker).unwrap_or_else(|| { + panic!( + "{} no longer declares `{marker}`: the block was renamed or reformatted, \ + update this parser and the matching one in DeltaScanContribSuite", + path.display() + ) + }) + marker.len(); + // No key literal contains a parenthesis, so the first `)` closes the `Seq(`. + let end = start + + scala[start..] + .find(')') + .expect("unterminated S3ConfigKeyConsumers Seq("); + let block = &scala[start..end]; + + let mut compared: BTreeSet<&str> = BTreeSet::new(); + let opening = "\"fs.s3a."; + let mut pos = 0usize; + while let Some(rel) = block[pos..].find(opening) { + let key_start = pos + rel + 1; + let key_end = key_start + + block[key_start..] + .find('"') + .expect("unterminated key literal in S3ConfigKeyConsumers"); + if block[key_end + 1..].trim_start().starts_with("->") { + compared.insert(&block[key_start..key_end]); + } + pos = key_end + 1; + } + assert!( + !compared.is_empty(), + "parsed no `\"fs.s3a.\" -> ` pairs out of S3ConfigKeyConsumers in {}", + path.display() + ); + + let missing: Vec = NATIVE_S3A_CONFIG_PROPERTIES + .iter() + .map(|p| format!("fs.s3a.{p}")) + .filter(|k| !compared.contains(k.as_str())) + .collect(); + assert!( + missing.is_empty(), + "native reads {missing:?} but S3ConfigKeyConsumers in {} does not compare them. Add \ + each with the tier matching how hadoop-aws reads it (LookupPasswordConsumer for \ + S3AUtils#lookupPassword credentials, PropagatedOptionConsumer for a plain \ + Configuration#get after propagateBucketOptions); a wrong tier compares the wrong \ + Hadoop value.", + path.display() + ); + } + /// Test configuration builder for easier setup Hadoop configurations #[derive(Debug, Default)] struct TestConfigBuilder { diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index fba2d6fc13e..688906abd09 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -23,7 +23,7 @@ use crate::parquet::parquet_support::{ object_store_authority, ObjectStoreBackend, SparkParquetOptions, }; use crate::parquet::schema_adapter::SparkPhysicalExprAdapterFactory; -use arrow::datatypes::{Field, FieldRef, SchemaRef}; +use arrow::datatypes::{Field, FieldRef, Schema, SchemaRef}; use datafusion::common::tree_node::TreeNode; use datafusion::config::{ParquetOptions, TableParquetOptions}; use datafusion::datasource::listing::PartitionedFile; @@ -47,6 +47,11 @@ use std::sync::Arc; #[cfg(test)] mod variant_tests; +/// Footer/page-index prefetch size for metadata reads, same as DataFusion's default. Shared +/// with the Delta DV path so its cache-populating footer fetch issues the identical read the +/// scan would. +pub(crate) const METADATA_SIZE_HINT: usize = 512 * 1024; + /// Initializes a DataSourceExec plan with a ParquetSource for Comet's native Parquet scan. /// /// `required_schema`: Schema to be projected by the scan. @@ -86,6 +91,9 @@ pub(crate) fn init_datasource_exec( encryption_enabled: bool, use_field_id: bool, require_field_ids: bool, + rebase_from_file_metadata: bool, + datetime_rebase_mode_in_read: &str, + int96_rebase_mode_in_read: &str, ) -> Result, ExecutionError> { // Computed once and reused below for `try_pushdown_filters`. `copied_config()` clones only // `SessionConfig` (an `Arc` plus a small extensions map); `SessionContext:: @@ -109,6 +117,9 @@ pub(crate) fn init_datasource_exec( // existing safe cast for filtered scans and use checked conversion only when every value is // necessarily read. spark_parquet_options.checked_timestamp_overflow = data_filters.is_none(); + spark_parquet_options.rebase_from_file_metadata = rebase_from_file_metadata; + spark_parquet_options.datetime_rebase_mode_in_read = datetime_rebase_mode_in_read.to_string(); + spark_parquet_options.int96_rebase_mode_in_read = int96_rebase_mode_in_read.to_string(); // Determine the schema and projection to use for ParquetSource. // When data_schema is provided, use it as the base schema so DataFusion knows the full @@ -140,6 +151,36 @@ pub(crate) fn init_datasource_exec( } _ => (Arc::clone(&required_schema), None), }; + + // DataFusion's parquet opener skips the physical-expr adapter entirely when no predicate + // is pushed down AND the logical and physical file schemas compare equal (the + // `needs_rewrite` fast path in `opener/mod.rs`). A parquet file with no footer key-value + // metadata -- exactly the non-Spark files whose rebase policy falls back to the session + // read modes -- can produce a physical schema identical to the logical one, silently + // bypassing the per-file rebase handling (which must refuse, or rebase, ancient values). + // Stamp a marker into the logical file schema's metadata so that equality can never hold + // for a rebase-enabled scan: parquet footers do not produce this key (Spark-written files + // carry `org.apache.spark.*` pairs that already break equality, and a crafted file + // embedding the marker via `ARROW:schema` merely degrades to the skip behavior). The + // marker propagates into `DataSourceExec::schema()`'s schema-level metadata (TableSchema + // copies it); that stays native-side only -- the JVM FFI export in `prepare_output` reads + // per-FIELD metadata, never the schema-level map -- but a future consumer comparing this + // scan's full `Schema` (metadata included) against an independently built one must expect + // the key. + let base_schema = if rebase_from_file_metadata { + let mut metadata = base_schema.metadata().clone(); + metadata.insert( + "comet.rebase_from_file_metadata".to_string(), + "true".to_string(), + ); + Arc::new(Schema::new_with_metadata( + base_schema.fields().clone(), + metadata, + )) + } else { + base_schema + }; + let partition_fields: Vec = partition_schema .iter() .flat_map(|s| s.fields().iter()) @@ -151,7 +192,7 @@ pub(crate) fn init_datasource_exec( let mut parquet_source = ParquetSource::new(table_schema) .with_table_parquet_options(table_parquet_options) - .with_metadata_size_hint(512 * 1024); // Same as DataFusion's default + .with_metadata_size_hint(METADATA_SIZE_HINT); let projects_variant = required_schema .fields() @@ -187,6 +228,10 @@ pub(crate) fn init_datasource_exec( let runtime_env = session_ctx.runtime_env(); let store = runtime_env.object_store(&object_store_url)?; let metadata_cache = runtime_env.cache_manager.get_file_metadata_cache(); + // + // A rebase-enabled scan also has the factory stamp each file's INT96 leaf ordinals into + // its footer metadata (see `datetime_rebase.rs`), which is how the expression adapter + // attributes timestamp columns to Spark's INT64 vs INT96 rebase specs. let scan_io_source = scan_io_source(object_store_backend); let reader_factory = Arc::new( EagerPageIndexReaderFactory::new( @@ -197,6 +242,7 @@ pub(crate) fn init_datasource_exec( ) .with_spark_variant_schema(projects_variant) .with_require_field_ids(require_field_ids) + .with_int96_leaf_stamp(rebase_from_file_metadata) .with_conversion_check( Arc::clone(&required_schema), spark_parquet_options.clone(), @@ -379,7 +425,7 @@ fn get_options( #[cfg(test)] mod tests { use super::*; - use arrow::array::Int32Array; + use arrow::array::{Date32Array, Int32Array}; use arrow::datatypes::{DataType, Field, Schema}; use arrow::record_batch::RecordBatch; use bytes::Bytes; @@ -498,6 +544,9 @@ mod tests { false, false, false, + false, + "", + "", ) .unwrap() } @@ -576,6 +625,572 @@ mod tests { } } + /// End-to-end pin for the per-file datetime rebase (see `datetime_rebase.rs`): the parquet + /// footer's Spark writer metadata must survive DataFusion's opener into the expr adapter, + /// and the resulting scan must return rebased dates -- but ONLY when the arm opted in. + /// The hybrid day count a legacy writer stores for Julian `1500-01-01` is numerically the + /// proleptic day of `1500-01-10` (-171655); rebasing restores proleptic `1500-01-01` + /// (-171664), the exact 9-day shift of the silent-corruption repro. + async fn scan_legacy_date_file(rebase_from_file_metadata: bool) -> Vec { + let schema = Arc::new(Schema::new(vec![Field::new("d", DataType::Date32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Date32Array::from(vec![-171655, 0]))], + ) + .unwrap(); + + let filename = get_temp_filename() + .as_path() + .as_os_str() + .to_str() + .unwrap() + .to_string(); + let props = WriterProperties::builder() + .set_key_value_metadata(Some(vec![ + KeyValue::new("org.apache.spark.version".to_string(), "3.5.9".to_string()), + KeyValue::new( + "org.apache.spark.legacyDateTime".to_string(), + "".to_string(), + ), + KeyValue::new("org.apache.spark.timeZone".to_string(), "UTC".to_string()), + ])) + .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 partitioned_file = PartitionedFile::from_path(filename).unwrap(); + let session_ctx = Arc::new(SessionContext::new()); + let scan = init_datasource_exec( + Arc::clone(&schema), + None, + None, + ObjectStoreUrl::local_filesystem(), + ObjectStoreBackend::Local, + vec![vec![partitioned_file]], + None, + None, + None, + "UTC", + true, + false, + false, + false, + &session_ctx, + false, + false, + false, + rebase_from_file_metadata, + "", + "", + ) + .unwrap(); + + let mut values = Vec::new(); + let mut stream = scan.execute(0, session_ctx.task_ctx()).unwrap(); + while let Some(batch) = stream.next().await { + let batch = batch.unwrap(); + let dates = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + values.extend(dates.iter().map(|v| v.unwrap())); + } + values + } + + #[tokio::test] + async fn rebases_legacy_dates_from_file_metadata_when_opted_in() { + assert_eq!(scan_legacy_date_file(true).await, vec![-171664, 0]); + } + + #[tokio::test] + async fn keeps_no_rebase_behavior_when_not_opted_in() { + // NativeScan's documented behavior (#5010): the legacy flag is ignored and the raw + // day count comes back unchanged. + assert_eq!(scan_legacy_date_file(false).await, vec![-171655, 0]); + } + + /// End-to-end pin for the session-read-mode fallback: a file with NO Spark writer metadata + /// (a non-Spark writer) resolves its rebase policy from the forwarded read modes -- + /// `DataSourceUtils.getRebaseSpec`'s `modeByConfig` fallback -- which must survive + /// `init_datasource_exec` into the expr adapter. + async fn scan_no_metadata_date_file(datetime_rebase_mode: &str) -> Result, String> { + let schema = Arc::new(Schema::new(vec![Field::new("d", DataType::Date32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Date32Array::from(vec![-171655, 0]))], + ) + .unwrap(); + + let filename = get_temp_filename() + .as_path() + .as_os_str() + .to_str() + .unwrap() + .to_string(); + let file = File::create(&filename).unwrap(); + let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let partitioned_file = PartitionedFile::from_path(filename).unwrap(); + let session_ctx = Arc::new(SessionContext::new()); + let scan = init_datasource_exec( + Arc::clone(&schema), + None, + None, + ObjectStoreUrl::local_filesystem(), + ObjectStoreBackend::Local, + vec![vec![partitioned_file]], + None, + None, + None, + "UTC", + true, + false, + false, + false, + &session_ctx, + false, + false, + false, + true, + datetime_rebase_mode, + datetime_rebase_mode, + ) + .unwrap(); + + let mut values = Vec::new(); + let mut stream = scan.execute(0, session_ctx.task_ctx()).unwrap(); + while let Some(batch) = stream.next().await { + let batch = batch.map_err(|e| e.to_string())?; + let dates = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + values.extend(dates.iter().map(|v| v.unwrap())); + } + Ok(values) + } + + #[tokio::test] + async fn non_spark_file_reads_ancient_dates_verbatim_under_corrected_read_mode() { + // Spark 4.0's default read mode: values pass through untouched, ancient included. + assert_eq!( + scan_no_metadata_date_file("CORRECTED").await.unwrap(), + vec![-171655, 0] + ); + } + + #[tokio::test] + async fn non_spark_file_rebases_ancient_dates_under_legacy_read_mode() { + // LEGACY read mode: the stored hybrid-calendar day count rebases to proleptic + // Gregorian (dates are zone-free, so the full rebase applies). + assert_eq!( + scan_no_metadata_date_file("LEGACY").await.unwrap(), + vec![-171664, 0] + ); + } + + #[tokio::test] + async fn non_spark_file_refuses_ancient_dates_under_default_read_mode() { + // An empty mode (older proto producer) keeps the conservative EXCEPTION posture. + let err = scan_no_metadata_date_file("").await.unwrap_err(); + assert!(err.contains("rebase"), "unexpected error: {err}"); + } + + /// Writes a metadata-free parquet file with one INT64 `TIMESTAMP_MICROS` column (`ts`) and + /// one INT96 column (`ts96`) through the low-level writer (arrow's writer cannot emit + /// INT96), then scans it with the given session read modes. `int96_days` is the day count + /// since the epoch the INT96 value nominally encodes (its Julian Day Number is + /// `2440588 + int96_days`). Returns the two values of the single row. + async fn scan_int96_and_int64_file( + int64_micros: i64, + int96_days: i32, + datetime_rebase_mode: &str, + int96_rebase_mode: &str, + ) -> Result<(i64, i64), String> { + use parquet::data_type::{Int64Type, Int96, Int96Type}; + use parquet::file::writer::SerializedFileWriter; + use parquet::schema::parser::parse_message_type; + + let filename = get_temp_filename() + .as_path() + .as_os_str() + .to_str() + .unwrap() + .to_string(); + let parquet_schema = Arc::new( + parse_message_type( + "message m { required int64 ts (TIMESTAMP(MICROS,true)); required int96 ts96; }", + ) + .unwrap(), + ); + let file = File::create(&filename).unwrap(); + let mut writer = + SerializedFileWriter::new(file, parquet_schema, Arc::new(WriterProperties::default())) + .unwrap(); + let mut row_group = writer.next_row_group().unwrap(); + let mut col = row_group.next_column().unwrap().unwrap(); + col.typed::() + .write_batch(&[int64_micros], None, None) + .unwrap(); + col.close().unwrap(); + let mut col = row_group.next_column().unwrap().unwrap(); + let mut int96 = Int96::new(); + int96.set_data(0, 0, (2_440_588 + int96_days as i64) as u32); + col.typed::() + .write_batch(&[int96], None, None) + .unwrap(); + col.close().unwrap(); + row_group.close().unwrap(); + writer.close().unwrap(); + + let ts_type = + DataType::Timestamp(arrow::datatypes::TimeUnit::Microsecond, Some("UTC".into())); + let schema = Arc::new(Schema::new(vec![ + Field::new("ts", ts_type.clone(), false), + Field::new("ts96", ts_type, false), + ])); + let partitioned_file = PartitionedFile::from_path(filename).unwrap(); + let session_ctx = Arc::new(SessionContext::new()); + let scan = init_datasource_exec( + Arc::clone(&schema), + None, + None, + ObjectStoreUrl::local_filesystem(), + ObjectStoreBackend::Local, + vec![vec![partitioned_file]], + None, + None, + None, + "UTC", + true, + false, + false, + false, + &session_ctx, + false, + false, + false, + true, + datetime_rebase_mode, + int96_rebase_mode, + ) + .unwrap(); + + let mut stream = scan.execute(0, session_ctx.task_ctx()).unwrap(); + let mut values = Vec::new(); + while let Some(batch) = stream.next().await { + let batch = batch.map_err(|e| e.to_string())?; + let ts = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let ts96 = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + values.extend( + ts.values() + .iter() + .zip(ts96.values().iter()) + .map(|(a, b)| (*a, *b)), + ); + } + assert_eq!(values.len(), 1); + Ok(values[0]) + } + + /// Proleptic `1500-01-01T00:00:00Z` in days / micros since the epoch. + const ANCIENT_DAYS: i32 = -171_664; + const ANCIENT_MICROS: i64 = ANCIENT_DAYS as i64 * 86_400_000_000; + + #[tokio::test] + async fn int64_timestamps_follow_the_datetime_spec_when_the_int96_spec_differs() { + // Spark selects `datetimeRebaseSpec` for INT64 MICROS/MILLIS columns and `int96RebaseSpec` + // only for INT96 columns: under datetime CORRECTED + int96 EXCEPTION, an ancient INT64 + // timestamp reads verbatim even though the INT96 spec would refuse an ancient INT96 + // value. The INT96 column here holds a modern value, so the whole row must read. + assert_eq!( + scan_int96_and_int64_file(ANCIENT_MICROS, 0, "CORRECTED", "EXCEPTION") + .await + .unwrap(), + (ANCIENT_MICROS, 0) + ); + } + + #[tokio::test] + async fn int96_timestamps_follow_the_int96_spec() { + // Same modes, ancient INT96 value: the INT96 spec (EXCEPTION) refuses it, naming the + // INT96 column -- not the INT64 one, which is fine under CORRECTED. + let err = scan_int96_and_int64_file(0, ANCIENT_DAYS, "CORRECTED", "EXCEPTION") + .await + .unwrap_err(); + assert!(err.contains("rebase"), "unexpected error: {err}"); + assert!(err.contains("'ts96'"), "unexpected error: {err}"); + + // Mirror image: datetime EXCEPTION + int96 CORRECTED reads an ancient INT96 value + // verbatim while a modern INT64 value passes the check. + assert_eq!( + scan_int96_and_int64_file(0, ANCIENT_DAYS, "EXCEPTION", "CORRECTED") + .await + .unwrap(), + (0, ANCIENT_MICROS) + ); + } + + /// The one value of a raw timestamp column, written through the low-level writer. + enum RawTimestamp { + Int64(i64), + /// Midnight of the day with this Julian Day Number, INT96-encoded. + Int96Midnight(u32), + } + + /// Writes a parquet file holding the single column `ts` of `message_type` (footer key-value + /// pairs from `key_values`, none by default: a non-Spark writer) with the one value + /// `value`, then scans it with `ts` requested as `requested` under the given session read + /// modes. Returns the column's microsecond values. + async fn scan_single_timestamp_column( + message_type: &str, + value: RawTimestamp, + key_values: Option>, + requested: DataType, + allow_timestamp_ltz_to_ntz: bool, + datetime_rebase_mode: &str, + int96_rebase_mode: &str, + ) -> Result, String> { + use parquet::data_type::{Int64Type, Int96, Int96Type}; + use parquet::file::writer::SerializedFileWriter; + use parquet::schema::parser::parse_message_type; + + let filename = get_temp_filename() + .as_path() + .as_os_str() + .to_str() + .unwrap() + .to_string(); + let parquet_schema = Arc::new(parse_message_type(message_type).unwrap()); + let props = WriterProperties::builder() + .set_key_value_metadata(key_values) + .build(); + let file = File::create(&filename).unwrap(); + let mut writer = SerializedFileWriter::new(file, parquet_schema, Arc::new(props)).unwrap(); + let mut row_group = writer.next_row_group().unwrap(); + let mut col = row_group.next_column().unwrap().unwrap(); + match value { + RawTimestamp::Int64(micros) => { + col.typed::() + .write_batch(&[micros], None, None) + .unwrap(); + } + RawTimestamp::Int96Midnight(julian_day) => { + let mut int96 = Int96::new(); + int96.set_data(0, 0, julian_day); + col.typed::() + .write_batch(&[int96], None, None) + .unwrap(); + } + } + col.close().unwrap(); + row_group.close().unwrap(); + writer.close().unwrap(); + + let schema = Arc::new(Schema::new(vec![Field::new("ts", requested, false)])); + let partitioned_file = PartitionedFile::from_path(filename).unwrap(); + let session_ctx = Arc::new(SessionContext::new()); + let scan = init_datasource_exec( + Arc::clone(&schema), + None, + None, + ObjectStoreUrl::local_filesystem(), + ObjectStoreBackend::Local, + vec![vec![partitioned_file]], + None, + None, + None, + "UTC", + true, + false, + false, + allow_timestamp_ltz_to_ntz, + &session_ctx, + false, + false, + false, + true, + datetime_rebase_mode, + int96_rebase_mode, + ) + .unwrap(); + + let mut stream = scan.execute(0, session_ctx.task_ctx()).unwrap(); + let mut values = Vec::new(); + while let Some(batch) = stream.next().await { + let batch = batch.map_err(|e| e.to_string())?; + let ts = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + values.extend(ts.values().iter().copied()); + } + Ok(values) + } + + /// Julian `1500-01-01T00:00:00` as a legacy writer stores it: the hybrid day count is + /// numerically the proleptic day of `1500-01-10`, so a UTC rebase restores `ANCIENT_MICROS`. + const HYBRID_1500_MICROS: i64 = -171_655 * 86_400_000_000; + /// Proleptic `1800-01-01T00:00:00Z` as a Julian Day Number and in micros since the epoch: + /// 62091 days before the epoch, ancient by Spark's 1900-01-01 timestamp cutoff. + const JDN_1800_01_01: u32 = 2_378_497; + const MICROS_1800_01_01: i64 = (JDN_1800_01_01 as i64 - 2_440_588) * 86_400_000_000; + + const TZ_FREE_INT64_MICROS: &str = "message m { required int64 ts (TIMESTAMP(MICROS,false)); }"; + const ADJUSTED_INT64_MICROS: &str = "message m { required int64 ts (TIMESTAMP(MICROS,true)); }"; + const INT96: &str = "message m { required int96 ts; }"; + + fn ltz_micros() -> DataType { + DataType::Timestamp(arrow::datatypes::TimeUnit::Microsecond, Some("UTC".into())) + } + + fn ntz_micros() -> DataType { + DataType::Timestamp(arrow::datatypes::TimeUnit::Microsecond, None) + } + + #[tokio::test] + async fn tz_free_int64_timestamps_requested_as_ltz_follow_the_datetime_spec() { + // Spark's INT64 branch keys on the requested type and checks only the unit + // (isTimestampTypeMatched), so a TIMESTAMP(MICROS, isAdjustedToUTC=false) column read + // as TIMESTAMP takes LongWithRebaseUpdater under the datetime spec: EXCEPTION refuses + // the ancient value, CORRECTED passes it verbatim, and LEGACY rebases it (exactly with + // a UTC writer zone, refused without one, since the JVM default zone is unknown here). + let err = scan_single_timestamp_column( + TZ_FREE_INT64_MICROS, + RawTimestamp::Int64(HYBRID_1500_MICROS), + None, + ltz_micros(), + false, + "EXCEPTION", + "CORRECTED", + ) + .await + .unwrap_err(); + assert!(err.contains("rebase"), "unexpected error: {err}"); + assert!(err.contains("'ts'"), "unexpected error: {err}"); + + assert_eq!( + scan_single_timestamp_column( + TZ_FREE_INT64_MICROS, + RawTimestamp::Int64(HYBRID_1500_MICROS), + None, + ltz_micros(), + false, + "CORRECTED", + "EXCEPTION", + ) + .await + .unwrap(), + vec![HYBRID_1500_MICROS] + ); + + let err = scan_single_timestamp_column( + TZ_FREE_INT64_MICROS, + RawTimestamp::Int64(HYBRID_1500_MICROS), + None, + ltz_micros(), + false, + "LEGACY", + "CORRECTED", + ) + .await + .unwrap_err(); + assert!(err.contains("timezone tables"), "unexpected error: {err}"); + + assert_eq!( + scan_single_timestamp_column( + TZ_FREE_INT64_MICROS, + RawTimestamp::Int64(HYBRID_1500_MICROS), + Some(vec![KeyValue::new( + "org.apache.spark.timeZone".to_string(), + "UTC".to_string(), + )]), + ltz_micros(), + false, + "LEGACY", + "CORRECTED", + ) + .await + .unwrap(), + vec![ANCIENT_MICROS] + ); + + // Read as TIMESTAMP_NTZ, the same column takes LongUpdater: no rebase in any mode. + assert_eq!( + scan_single_timestamp_column( + TZ_FREE_INT64_MICROS, + RawTimestamp::Int64(HYBRID_1500_MICROS), + None, + ntz_micros(), + false, + "EXCEPTION", + "EXCEPTION", + ) + .await + .unwrap(), + vec![HYBRID_1500_MICROS] + ); + } + + #[tokio::test] + async fn timestamps_requested_as_ntz_are_never_rebased() { + // Spark 4.x reads an INT96 column as TIMESTAMP_NTZ through BinaryToSQLTimestampUpdater + // and an adjusted INT64 column through LongUpdater; neither consults a rebase mode + // (Spark 3.x refuses the pairing before any rebase decision, which Comet's + // allow_timestamp_ltz_to_ntz gate reproduces). + for (datetime_mode, int96_mode) in [ + ("EXCEPTION", "EXCEPTION"), + ("CORRECTED", "CORRECTED"), + ("LEGACY", "LEGACY"), + ] { + assert_eq!( + scan_single_timestamp_column( + INT96, + RawTimestamp::Int96Midnight(JDN_1800_01_01), + None, + ntz_micros(), + true, + datetime_mode, + int96_mode, + ) + .await + .unwrap_or_else(|e| panic!("{datetime_mode}/{int96_mode}: {e}")), + vec![MICROS_1800_01_01], + "{datetime_mode}/{int96_mode}" + ); + } + assert_eq!( + scan_single_timestamp_column( + ADJUSTED_INT64_MICROS, + RawTimestamp::Int64(ANCIENT_MICROS), + None, + ntz_micros(), + true, + "EXCEPTION", + "EXCEPTION", + ) + .await + .unwrap(), + vec![ANCIENT_MICROS] + ); + } + // Regression test for #4990: a fresh `TableParquetOptions::new()` ignored session-level // `datafusion.execution.parquet.*` settings entirely, so `spark.comet.datafusion. // execution.parquet.*` (behind `respectDataFusionConfigs`) and `spark.comet.parquet. @@ -678,6 +1293,9 @@ mod tests { false, false, false, + false, + "", + "", ) .unwrap(); diff --git a/native/core/src/parquet/parquet_exec/variant_tests.rs b/native/core/src/parquet/parquet_exec/variant_tests.rs index ec857019656..0bfdcecd222 100644 --- a/native/core/src/parquet/parquet_exec/variant_tests.rs +++ b/native/core/src/parquet/parquet_exec/variant_tests.rs @@ -163,6 +163,9 @@ async fn scan_variant_file(filename: PathBuf) -> VariantArray { false, false, false, + false, + "", + "", ) .unwrap(); let mut stream = scan.execute(0, session_ctx.task_ctx()).unwrap(); @@ -228,6 +231,9 @@ async fn unread_variant_does_not_override_arrow_schema_hint() { false, false, false, + false, + "", + "", ) .unwrap(); let mut stream = scan.execute(0, session.task_ctx()).unwrap(); @@ -258,6 +264,9 @@ fn encrypted_projected_variant_is_rejected_before_reader_creation() { true, false, false, + false, + "", + "", ); assert!(result .unwrap_err() diff --git a/native/core/src/parquet/parquet_support.rs b/native/core/src/parquet/parquet_support.rs index 8fc807a14cb..8d10f80d4c2 100644 --- a/native/core/src/parquet/parquet_support.rs +++ b/native/core/src/parquet/parquet_support.rs @@ -116,6 +116,22 @@ pub struct SparkParquetOptions { /// (overflow -> NULL), because Spark may discard values through pruning paths that /// DataFusion cannot fully mirror before conversion. pub checked_timestamp_overflow: bool, + /// When true, resolve each file's datetime calendar-rebase policy from its parquet footer + /// metadata (`org.apache.spark.legacyDateTime` and friends) and rebase -- or refuse -- + /// affected values, mirroring Spark's per-file `DataSourceUtils.datetimeRebaseSpec` + /// resolution. Enabled by the Delta scan arm; the plain NativeScan keeps its documented + /// no-rebase behavior (#5010). See `datetime_rebase.rs`. + pub rebase_from_file_metadata: bool, + /// Effective `spark.sql.parquet.datetimeRebaseModeInRead` (a `LegacyBehaviorPolicy` value), + /// forwarded from the JVM at planning time. Consulted -- exactly like Spark's + /// `DataSourceUtils.getRebaseSpec` `modeByConfig` fallback -- only for files whose footer + /// metadata does not decide the rebase policy on its own, and only when + /// `rebase_from_file_metadata` is set. Empty (a producer that predates the field) is + /// treated as `EXCEPTION`, the conservative refuse-ancient posture. + pub datetime_rebase_mode_in_read: String, + /// Effective `spark.sql.parquet.int96RebaseModeInRead`; same semantics as + /// `datetime_rebase_mode_in_read` for the INT96 timestamp spec. + pub int96_rebase_mode_in_read: String, } impl SparkParquetOptions { @@ -131,6 +147,9 @@ impl SparkParquetOptions { allow_type_promotion: false, allow_timestamp_ltz_to_ntz: false, checked_timestamp_overflow: true, + rebase_from_file_metadata: false, + datetime_rebase_mode_in_read: String::new(), + int96_rebase_mode_in_read: String::new(), } } @@ -146,6 +165,9 @@ impl SparkParquetOptions { allow_type_promotion: false, allow_timestamp_ltz_to_ntz: false, checked_timestamp_overflow: true, + rebase_from_file_metadata: false, + datetime_rebase_mode_in_read: String::new(), + int96_rebase_mode_in_read: String::new(), } } } @@ -675,7 +697,7 @@ pub(crate) fn object_store_authority(url: &Url) -> &str { &url[start..url::Position::AfterPort] } -fn object_store_url_key(url: &Url) -> String { +fn registry_url_key(url: &Url) -> String { format!("{}://{}", url.scheme(), object_store_authority(url)) } @@ -698,7 +720,7 @@ impl ObjectStoreRegistry for CometObjectStoreRegistry { self.azure_stores .write() .unwrap_or_else(PoisonError::into_inner) - .insert(object_store_url_key(url), store) + .insert(registry_url_key(url), store) } else { self.default.register_store(url, store) } @@ -709,7 +731,7 @@ impl ObjectStoreRegistry for CometObjectStoreRegistry { self.azure_stores .write() .unwrap_or_else(PoisonError::into_inner) - .remove(&object_store_url_key(url)) + .remove(®istry_url_key(url)) .ok_or_else(|| { DataFusionError::Internal(format!("No suitable object store found for {url}")) }) @@ -724,7 +746,7 @@ impl ObjectStoreRegistry for CometObjectStoreRegistry { .azure_stores .read() .unwrap_or_else(PoisonError::into_inner) - .get(&object_store_url_key(url)) + .get(®istry_url_key(url)) { return Ok(Arc::clone(store)); } @@ -837,13 +859,13 @@ type ObjectStoreCache = RwLock /// (e.g. `fs.s3a.access.key` / `fs.s3a.secret.key`) produce a different `config_hash` when /// those values change, which causes a new store to be created and inserted under the new /// key; the old entry is harmlessly superseded. -fn object_store_cache() -> &'static ObjectStoreCache { +pub(crate) fn object_store_cache() -> &'static ObjectStoreCache { static CACHE: OnceLock = OnceLock::new(); CACHE.get_or_init(|| RwLock::new(HashMap::new())) } /// Compute a hash of the object store configuration for cache keying. -fn hash_object_store_configs(configs: &HashMap) -> u64 { +pub(crate) fn hash_object_store_configs(configs: &HashMap) -> u64 { let mut hasher = DefaultHasher::new(); let mut keys: Vec<&String> = configs.keys().collect(); keys.sort(); @@ -889,25 +911,52 @@ fn object_store_backend(url: &Url, is_hdfs: bool) -> Result, url: String, object_store_configs: &HashMap, +) -> Result<(ObjectStoreUrl, Path, ObjectStoreBackend), ExecutionError> { + let config_hash = hash_object_store_configs(object_store_configs); + prepare_object_store_with_config_hash(runtime_env, url, object_store_configs, config_hash) +} + +/// The `scheme://host:port` cache-key string [`prepare_object_store_with_configs`] resolves and +/// registers object stores under, plus the "is this an HDFS-scheme URL" classification. `url` +/// must already be the [`normalize_object_store_url`] result (s3a and the opted-in aliases +/// rewritten to `s3://`, a hostless alias bucket promoted into the host), so the key here is +/// exactly the one the resolution path registers under. Pure and I/O-free (no config hashing, +/// no cache lock, no store creation/registration): a caller that keeps its OWN local +/// `ObjectStoreUrl`-keyed cache (e.g. `delta_spark_scan.rs`'s `resolve_store`, which resolves a +/// store per FILE but only needs one per distinct authority) can compute this cheap key first +/// and consult its local cache before ever calling into the expensive resolution path below. +pub(crate) fn object_store_url_key(normalized: &NormalizedObjectStoreUrl) -> (String, bool) { + let url = &normalized.url; + let url_key = format!("{}://{}", url.scheme(), object_store_authority(url)); + (url_key, normalized.is_hdfs) +} + +/// Same as [`prepare_object_store_with_configs`], but takes an already-computed +/// [`hash_object_store_configs`] result instead of hashing `object_store_configs` again. `configs` +/// is loop-invariant across every file resolved for one scan/writer, so a caller that already +/// hashed it once (e.g. once per partition, rather than once per file) should call this directly. +pub(crate) fn prepare_object_store_with_config_hash( + runtime_env: Arc, + url: String, + object_store_configs: &HashMap, + config_hash: u64, ) -> Result<(ObjectStoreUrl, Path, ObjectStoreBackend), ExecutionError> { // `is_hdfs` comes back from normalization because it must be decided on the URL as written. // Re-deriving it from the normalized URL would let an `s3a`/alias rewrite land on an `s3` // entry in `fs.comet.libhdfs.schemes` and route an S3 read through libhdfs. - let NormalizedObjectStoreUrl { - url, - is_hdfs: is_hdfs_scheme, - } = normalize_object_store_url(url.as_str(), object_store_configs)?; + let normalized = normalize_object_store_url(url.as_str(), object_store_configs)?; + let (url_key, is_hdfs_scheme) = object_store_url_key(&normalized); // Configured S3 aliases must be normalized before the object-store parser classifies them. // HDFS routing still wins, including when its configured schemes resemble remote stores. - let backend = object_store_backend(&url, is_hdfs_scheme)?; - let scheme = url.scheme(); - let url_key = object_store_url_key(&url); + let backend = object_store_backend(&normalized.url, is_hdfs_scheme)?; + let url = &normalized.url; - let config_hash = hash_object_store_configs(object_store_configs); let cache_key = (url_key.clone(), config_hash, is_hdfs_scheme); // Check the cache first to reuse existing object store instances. @@ -928,13 +977,13 @@ pub(crate) fn prepare_object_store_with_configs( } else { debug!("Creating new object store for {url_key}"); let (store, path): (Box, Path) = if is_hdfs_scheme { - create_hdfs_object_store(&url) - } else if scheme == "s3" { - objectstore::s3::create_store(&url, object_store_configs, Duration::from_secs(300)) - } else if is_azure_scheme(scheme) { - objectstore::azure::create_store(&url, object_store_configs) + create_hdfs_object_store(url) + } else if url.scheme() == "s3" { + objectstore::s3::create_store(url, object_store_configs, Duration::from_secs(300)) + } else if is_azure_scheme(url.scheme()) { + objectstore::azure::create_store(url, object_store_configs) } else { - parse_url(&url) + parse_url(url) } .map_err(|e| ExecutionError::GeneralError(e.to_string()))?; @@ -946,31 +995,40 @@ pub(crate) fn prepare_object_store_with_configs( (store, path) }; - // A RuntimeEnv can plan multiple scans with different backends or credentials - // for the same bucket. Use the same identity as the cache, even for the first - // registration, so neither later registration nor planning order changes the - // store used by an existing scan. Native s3/s3a share the normalized s3 scheme; - // a Hadoop-selected scheme retains its physical spelling. - // - // Native LocalFileSystem ignores these Hadoop options and keeps file:// for - // compatibility. An explicitly Hadoop-routed file scheme is still isolated. - let object_store_url = if scheme == "file" && !is_hdfs_scheme { - ObjectStoreUrl::parse(url_key)? - } else { - let backend = if is_hdfs_scheme { "hdfs" } else { "native" }; - // DataFusion keys stores only by scheme and authority, so put configuration - // and backend identity in the scheme while preserving the physical authority. - // `+comet-` marks our internal registration suffix; encryption lookup strips - // the complete suffix to recover the physical URI. - ObjectStoreUrl::parse(format!( - "{scheme}+comet-{config_hash:016x}-{backend}://{}", - object_store_authority(&url), - ))? - }; + // A RuntimeEnv can plan multiple scans with different backends or credentials for the + // same bucket. Register under the same identity as the cache, even the first time, so + // neither later registration nor planning order changes the store an existing scan uses. + let object_store_url = object_store_registration_url(&normalized, &url_key, config_hash)?; runtime_env.register_object_store(object_store_url.as_ref(), object_store); Ok((object_store_url, object_store_path, backend)) } +/// The URL [`prepare_object_store_with_config_hash`] registers `normalized` under in a +/// `RuntimeEnv`, given its [`object_store_url_key`] and [`hash_object_store_configs`] result. +/// Native `file` keeps the physical key (LocalFileSystem ignores the Hadoop options); every +/// other store, including a Hadoop-routed `file`, folds the configuration hash and backend into +/// the scheme because DataFusion keys stores only by scheme and authority. `+comet-` marks the +/// suffix encryption lookup strips to recover the physical URI. Pure and I/O-free, so a caller +/// memoizing stores per registration URL can compute the key without resolving anything. +pub(crate) fn object_store_registration_url( + normalized: &NormalizedObjectStoreUrl, + url_key: &str, + config_hash: u64, +) -> Result { + let url = &normalized.url; + if url.scheme() == "file" && !normalized.is_hdfs { + return Ok(ObjectStoreUrl::parse(url_key)?); + } + let backend = if normalized.is_hdfs { "hdfs" } else { "native" }; + // DataFusion keys stores only by scheme and authority, so put configuration and backend + // identity in the scheme while preserving the physical authority, container included. + Ok(ObjectStoreUrl::parse(format!( + "{}+comet-{config_hash:016x}-{backend}://{}", + url.scheme(), + object_store_authority(url), + ))?) +} + #[cfg(test)] mod tests { /// Checks parser-backed I/O labels without constructing stores, including libhdfs overrides @@ -1293,6 +1351,67 @@ mod tests { object_store_cache().write().unwrap().remove(&key); } + /// Guards `object_store_registration_url` against drifting from the URL the resolution path + /// actually registers: the Delta scan memoizes stores under the former and reads them back + /// under the latter. Seeds one in-memory hdfs cache entry so no libhdfs backend is built. + #[test] + #[cfg_attr(miri, ignore)] // AWS credential providers and object_store call foreign functions + fn registration_url_matches_prepare_for_every_backend() { + use super::{ + object_store_registration_url, object_store_url_key, + prepare_object_store_with_config_hash, + }; + use crate::parquet::objectstore::s3_blob_fs_support::normalize_object_store_url; + let s3_options = HashMap::from([ + ( + "fs.s3a.aws.credentials.provider".to_string(), + "org.apache.hadoop.fs.s3a.AnonymousAWSCredentialsProvider".to_string(), + ), + ( + "fs.s3a.endpoint.region".to_string(), + "us-east-1".to_string(), + ), + ]); + let hdfs_options = + HashMap::from([("fs.comet.libhdfs.schemes".to_string(), "hdfs".to_string())]); + let hdfs_key = ( + "hdfs://comet-registration-url:8020".to_string(), + hash_object_store_configs(&hdfs_options), + true, + ); + let hdfs_store: Arc = Arc::new(InMemory::new()); + object_store_cache() + .write() + .unwrap() + .insert(hdfs_key.clone(), hdfs_store); + for (input, options) in [ + ("s3a://comet-registration-url/a.parquet", &s3_options), + ( + "file:///tmp/comet-registration-url/a.parquet", + &HashMap::new(), + ), + ( + "hdfs://comet-registration-url:8020/a.parquet", + &hdfs_options, + ), + ] { + let config_hash = hash_object_store_configs(options); + let normalized = normalize_object_store_url(input, options).unwrap(); + let (url_key, _) = object_store_url_key(&normalized); + let expected = object_store_registration_url(&normalized, &url_key, config_hash) + .unwrap_or_else(|e| panic!("{input}: {e}")); + let (registered, _, _) = prepare_object_store_with_config_hash( + Arc::new(RuntimeEnv::default()), + input.to_string(), + options, + config_hash, + ) + .unwrap_or_else(|e| panic!("{input}: {e}")); + assert_eq!(registered, expected, "{input}"); + } + object_store_cache().write().unwrap().remove(&hdfs_key); + } + /// Checks that native file construction returns Local and cached Hadoop file routing returns /// Other, with distinct registered stores. Removes its synthetic Hadoop cache entry on success. #[test] diff --git a/native/core/src/parquet/schema_adapter.rs b/native/core/src/parquet/schema_adapter.rs index fb37111ec5f..5f3a3571945 100644 --- a/native/core/src/parquet/schema_adapter.rs +++ b/native/core/src/parquet/schema_adapter.rs @@ -16,6 +16,10 @@ // under the License. use crate::parquet::cast_column::CometCastColumnExpr; +use crate::parquet::datetime_rebase::{ + resolve_file_rebase_policies, wrap_datetime_rebase, FileRebasePolicies, RebaseReadMode, + SessionRebaseModes, +}; use crate::parquet::name_fold::{fold_name, fold_names, fold_schema_names}; use crate::parquet::parquet_support::{ duplicate_parquet_field_error, field_id, field_names_with_id, match_struct_fields, @@ -1015,6 +1019,58 @@ impl PhysicalExprAdapterFactory for SparkPhysicalExprAdapterFactory { Arc::clone(&adapted_physical_schema), )?; + // Per-file calendar-rebase policies, resolved from the ORIGINAL physical file schema: + // its metadata carries the parquet footer's key-value pairs (they survive the parquet + // -> arrow schema conversion; the remapped schema above rebuilds fields only and keeps + // no metadata) including the reader factory's INT96 leaf stamp, and its field tree + // validates that stamp. `None` -- the overwhelmingly common case -- means no wrapping + // in `rewrite` at all. + let rebase_policies = if self.parquet_options.rebase_from_file_metadata { + // The session read modes only matter for files without Spark writer metadata + // (getRebaseSpec's modeByConfig fallback); empty strings parse to EXCEPTION. + let session_modes = SessionRebaseModes { + datetime: RebaseReadMode::from_conf_value( + &self.parquet_options.datetime_rebase_mode_in_read, + ), + int96: RebaseReadMode::from_conf_value( + &self.parquet_options.int96_rebase_mode_in_read, + ), + }; + let policies = resolve_file_rebase_policies(&physical_file_schema, session_modes); + policies.any_rebase_needed().then(|| { + // Pair each physical column with the logical field the adapter narrows it to, + // through the folded names computed above (the remap already renamed id-matched + // columns to their logical names), so the wrapper -- which sits beneath the + // nested narrowing -- never checks a nested leaf the narrowing drops, and reads + // each timestamp leaf under the policy of its REQUESTED type (TIMESTAMP_NTZ + // never rebases, TIMESTAMP always takes a spec), as Spark's updater factory + // does. An unpaired physical column keeps every leaf under its physical type's + // policy; nothing references it anyway. + // First match wins on a folded-name collision, the same tie-break as + // `wrap_all_type_mismatches` and `remap_physical_schema`. + let mut logical_index: HashMap<&str, usize> = HashMap::new(); + for (i, name) in logical_folded.iter().enumerate() { + logical_index.entry(name.as_str()).or_insert(i); + } + let requested: Vec> = physical_folded + .iter() + .map(|name| { + logical_index + .get(name.as_str()) + .map(|&i| logical_file_schema.field(i).data_type()) + }) + .collect(); + policies.restrict_to_requested( + &physical_file_schema, + &requested, + case_sensitive, + self.parquet_options.use_field_id, + ) + }) + } else { + None + }; + Ok(Arc::new(SparkPhysicalExprAdapter { logical_file_schema, physical_file_schema: adapted_physical_schema, @@ -1027,6 +1083,7 @@ impl PhysicalExprAdapterFactory for SparkPhysicalExprAdapterFactory { id_duplicate_roots, logical_folded, physical_folded, + rebase_policies, })) } } @@ -1085,6 +1142,10 @@ struct SparkPhysicalExprAdapter { /// `physical_file_schema` field names pre-folded once, parallel to /// `physical_file_schema.fields()`. See `logical_folded`. physical_folded: Vec, + /// This file's datetime calendar-rebase policies, resolved once in `create` from the file's + /// footer metadata. `Some` only when `rebase_from_file_metadata` is set AND some policy is + /// not the plain proleptic-Gregorian pass-through; see `datetime_rebase.rs`. + rebase_policies: Option, } impl PhysicalExprAdapter for SparkPhysicalExprAdapter { @@ -1176,6 +1237,16 @@ impl PhysicalExprAdapter for SparkPhysicalExprAdapter { expr }; + // Last, wrap column references to this file's date/timestamp columns per its resolved + // calendar-rebase policies (Delta arm only; see `datetime_rebase.rs`). Runs after every + // remap so the wrap keys on the FINAL physical column indices, and wraps the raw column + // BENEATH any cast the adapters inserted, so casts see rebased (proleptic) values. + let expr = if let Some(policies) = &self.rebase_policies { + wrap_datetime_rebase(expr, &self.physical_file_schema, policies)? + } else { + expr + }; + Ok(expr) } } diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 97c65f0da36..f9292cb3f37 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -201,6 +201,76 @@ message NativeScan { optional int32 source_key_hash = 3; } +// Delta-table-wide data shared by all partitions (sent once at planning). +// Produced by the contrib Delta module; the native handler is compiled only +// when the `delta` Cargo feature is enabled. +message DeltaSparkScanCommon { + // Table root URL, used to resolve relative deletion-vector paths. + string table_root = 1; + // Column mapping mode: "none", "name", or "id". + string column_mapping_mode = 2; + // Key for split-mode plan-data injection. Derived from (table root, snapshot + // version, scan hash) so two scans of the same table in one plan (self-join, + // MERGE) don't collide -- same lesson as IcebergScan's + // (metadata_location, scan_hash_code) key. + string source_key = 3; + // Effective datetime rebase read modes (LegacyBehaviorPolicy values of + // spark.sql.parquet.datetimeRebaseModeInRead / int96RebaseModeInRead, resolved + // through ParquetOptions so per-relation options win, exactly as + // ParquetFileFormat.buildReaderWithPartitionValues resolves them). Consulted + // only for files whose footer metadata does not decide the rebase policy on + // its own (no org.apache.spark.version key), mirroring + // DataSourceUtils.getRebaseSpec's modeByConfig fallback. Empty (an older + // producer) is read as EXCEPTION, the conservative refuse-ancient posture. + string datetime_rebase_mode_in_read = 4; + string int96_rebase_mode_in_read = 5; +} + +// Descriptor for a Delta deletion vector, derived from the Delta protocol's +// DeletionVectorDescriptor. The JVM side (which has delta-spark on the +// classpath) resolves UUID-relative paths to absolute URLs and Z85-decodes +// inline bitmaps, so the native side needs neither codec. Executors fetch +// on-disk bitmaps with a single ranged object-store read; only this small +// descriptor crosses JNI. +message DeltaSparkDvDescriptor { + // Original storage form, for diagnostics: "u" (UUID-relative), "i" + // (inline), "p" (absolute path). + string storage_type = 1; + // Absolute URL of the DV file (on-disk forms). At descriptor.offset the + // file holds [i32 BE size][bitmap data][i32 BE CRC32-of-data]. + optional string absolute_path = 2; + // The bitmap data (magic + RoaringBitmapArray), already unframed and + // Z85-decoded (inline form). + optional bytes inline_data = 3; + // Byte offset of the size-prefixed bitmap within the DV file. + optional int32 offset = 4; + // Length of the bitmap data (excluding the size/CRC framing). + int32 size_in_bytes = 5; + // Number of deleted rows encoded in the bitmap. + int64 cardinality = 6; +} + +// A data file plus its optional deletion vector. +message DeltaSparkPartitionedFile { + SparkPartitionedFile file = 1; + optional DeltaSparkDvDescriptor dv = 2; +} + +// Single partition's Delta file list (injected at execution time). +// Field name matches SparkFilePartition.partitioned_file for consistency. +message DeltaSparkFilePartition { + repeated DeltaSparkPartitionedFile partitioned_file = 1; +} + +message DeltaSparkScan { + // Reuses the parquet scan's common data (schemas, filters, projections, + // object-store options, reader flags) -- the Delta read path delegates to + // the same native parquet machinery as NativeScan. + NativeScanCommon common = 1; + DeltaSparkScanCommon delta_common = 2; + DeltaSparkFilePartition file_partition = 3; +} + message CsvScan { repeated SparkStructField data_schema = 1; repeated SparkStructField partition_schema = 2; diff --git a/pom.xml b/pom.xml index 084c086e447..2a737c9c20c 100644 --- a/pom.xml +++ b/pom.xml @@ -45,11 +45,14 @@ under the License. UTF-8 17 4.1.0 @@ -105,6 +108,16 @@ under the License. Hadoop 3.4.x's CredentialProviderListFactory); the spark-3.4 / spark-3.5 profiles override it down to 3.3.x (SDK v1). hadoop.version itself (hadoop-client) is not changed. --> 3.4.1 + + delta-spark + 4.3.1 + delta-core + 2.4.0 17 ${java.version} ${java.version} @@ -743,6 +761,7 @@ under the License. spark-3.x spark-3.5 spark-none + 3.3.2 17 ${java.version} ${java.version} @@ -763,6 +782,7 @@ under the License. spark-4.x spark-4.0 spark-none + 4.0.1 17 ${java.version} ${java.version} @@ -786,6 +806,9 @@ under the License. spark-4.x spark-4.1+ spark-4.1 + + 4.3.1 17 ${java.version} ${java.version} @@ -807,6 +830,10 @@ under the License. spark-4.x spark-4.1+ spark-4.2 + + 4.3.1 17 ${java.version} @@ -814,6 +841,16 @@ under the License. + + + delta + + contrib/delta-spark + + + scala-2.12 diff --git a/spark/pom.xml b/spark/pom.xml index 6b71b199eb8..b6306449527 100644 --- a/spark/pom.xml +++ b/spark/pom.xml @@ -283,6 +283,27 @@ under the License. + + + delta + + + + org.apache.maven.plugins + maven-jar-plugin + + + + test-jar + + + + + + + celeborn-reflection-compatibility diff --git a/spark/src/main/scala/org/apache/comet/ContribServices.scala b/spark/src/main/scala/org/apache/comet/ContribServices.scala index 4cc0e2fa026..879ed2baac2 100644 --- a/spark/src/main/scala/org/apache/comet/ContribServices.scala +++ b/spark/src/main/scala/org/apache/comet/ContribServices.scala @@ -95,7 +95,7 @@ object ContribServices extends Logging { if (iterator.hasNext) found += iterator.next() else done = true } catch { - // NonFatal covers ServiceConfigurationError; LinkageError/OOM still propagate. + // NonFatal covers ServiceConfigurationError; OOM and the like still propagate. case NonFatal(e) => logWarning( s"Skipping an unusable ${service.getSimpleName} provider; the remaining providers " + @@ -103,6 +103,15 @@ object ContribServices extends Logging { "names a class that is absent, does not implement the service, or cannot be " + "constructed).", e) + // ServiceLoader raises NoClassDefFoundError while loading a provider whose superclass + // or interface is missing, the version-skewed-jar case. Contain it like a NonFatal + // failure so the remaining providers and every scan still work. + case e: LinkageError => + logWarning( + s"Skipping a ${service.getSimpleName} provider that cannot link " + + s"(${e.getClass.getName}); the remaining providers are unaffected. This usually " + + "means a contrib jar was built against a different Comet or Spark version.", + e) } } if (steps >= MaxSteps) { diff --git a/spark/src/main/scala/org/apache/comet/rules/CometScanContrib.scala b/spark/src/main/scala/org/apache/comet/rules/CometScanContrib.scala index 5bff024d00b..5a11e2dcc39 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometScanContrib.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometScanContrib.scala @@ -120,8 +120,15 @@ object CometScanContrib extends Logging { * speculative and can fail for reasons entirely outside the query -- an unreachable object * store, a metadata format newer than the contrib understands, a version-skewed reflective * lookup -- and none of those should turn a runnable query into a failed one. Logging (rather - * than swallowing silently) keeps an unexpectedly-declining contrib diagnosable. `NonFatal` - * deliberately lets `LinkageError`/`OOM`-class failures through. + * than swallowing silently) keeps an unexpectedly-declining contrib diagnosable. + * + * `NonFatal` does not match `LinkageError` (`NoSuchMethodError`, `NoClassDefFoundError`, ...), + * so it is caught separately and contained the same way: a contrib jar built against internals + * Comet has since moved or removed is a classpath/version skew, not a JVM-corrupting failure, + * and must not fail a query Spark could otherwise run. Genuinely fatal conditions -- + * `OutOfMemoryError` and the like -- are neither `NonFatal` nor `LinkageError` and always + * propagate; this is a narrow, deliberate widening for one specific `Error` subtype, not a + * blanket `catch (Throwable)`. */ private def firstClaim(hook: CometScanContrib => Option[SparkPlan]): Option[SparkPlan] = firstClaimFrom(contribs)(hook) @@ -147,6 +154,21 @@ object CometScanContrib extends Logging { "declining it and continuing with Comet's built-in handling", e) None + case e: LinkageError => + // A version-skewed contrib jar (compiled against a Comet internal that has since + // moved, been renamed, or been removed) surfaces as NoSuchMethodError, + // NoClassDefFoundError, or a sibling LinkageError -- a classpath mismatch, not a + // query-specific failure, and not the JVM corruption OutOfMemoryError/StackOverflowError + // signal. Contained the same way a NonFatal decline is: logged and treated as "this + // contrib does not claim this scan" so a stale contrib jar cannot fail a query Spark + // could otherwise run. + logWarning( + s"Contrib scan handler ${contrib.getClass.getName} failed with " + + s"${e.getClass.getName}, indicating it was built against a different version of " + + "Comet's internals than is on the classpath now; declining it and continuing with " + + "Comet's built-in handling", + e) + None } // Short-circuit before reading the config: a default build registers nothing, and this is on 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 edda7d3fabd..70e32bbee81 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala @@ -1223,7 +1223,8 @@ object CometScanRule extends Logging { * native can't be consulted (library not loaded), assume supported -- the gate is only an * early-fallback optimization and such a build can't run the native scan anyway. */ - private[rules] def isNativelyReadableScheme( + // private[comet] (not [rules]) so contrib scan extensions can apply the same gate. + private[comet] def isNativelyReadableScheme( uri: URI, s3CompliantSchemes: Set[String]): Boolean = { val scheme = uri.getScheme @@ -1290,7 +1291,8 @@ object CometScanRule extends Logging { catch { case _: Throwable => true } /** [[probeObjectStore]] against a URI's real path, not just its scheme. Uncached. */ - private[rules] def objectStoreAcceptsPath(uri: URI): Boolean = + // private[comet] (not [rules]) so contrib scan extensions can apply the same gate. + private[comet] def objectStoreAcceptsPath(uri: URI): Boolean = probeObjectStore(uri.toString) /** diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala index 38e7cdbf2e4..80324e62c7f 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala @@ -22,6 +22,7 @@ package org.apache.comet.serde.operator import scala.collection.mutable.ListBuffer import scala.jdk.CollectionConverters._ +import org.apache.hadoop.conf.Configuration import org.apache.spark.internal.Logging import org.apache.spark.sql.catalyst.analysis.Resolver import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, Expression, Literal} @@ -240,153 +241,218 @@ object CometNativeScan extends CometOperatorSerde[CometScanExec] with CometTypeS scan: CometScanExec, builder: Operator.Builder, childOp: OperatorOuterClass.Operator*): Option[OperatorOuterClass.Operator] = { - val nativeScanBuilder = OperatorOuterClass.NativeScan.newBuilder() + // Collect S3/cloud storage configurations + val hadoopConf = scan.relation.sparkSession.sessionState + .newHadoopConfWithOptions(scan.relation.options) + + // Use selectedPartitions (static) instead of getFilePartitions() because at planning time + // DPP subqueries haven't been resolved yet. Object store options don't depend on DPP. + val fileUris = scan.selectedPartitions.view.flatMap(_.files).map(_.getPath.toUri) + val objectStoreOptions = + NativeConfig.extractScanObjectStoreOptions(hadoopConf, fileUris) match { + case Right(options) => options + case Left(reason) => + // CometExecRule falls back to the wrapped Spark scan, so it carries the reason too. + withFallbackReason(scan, reason) + withFallbackReason(scan.wrapped, reason) + return None + } + + buildNativeScanCommon( + source = scan.simpleStringWithNodeId(), + output = scan.output, + requiredSchema = scan.requiredSchema, + dataSchema = scan.relation.dataSchema, + partitionSchema = scan.relation.partitionSchema, + fileConstantMetadataColumns = scan.wrapped.fileConstantMetadataColumns, + dataFilters = scan.supportedDataFilters, + objectStoreOptions = objectStoreOptions, + hadoopConf = hadoopConf, + conf = scan.conf) match { + case Some(commonBuilder) => + // Sink operators don't have children + builder.clearChildren() + val nativeScanBuilder = OperatorOuterClass.NativeScan.newBuilder() + // Set common data in NativeScan (file_partition will be populated at execution time) + nativeScanBuilder.setCommon(commonBuilder.build()) + Some(builder.setNativeScan(nativeScanBuilder).build()) + case None => + if (scan.output.forall(attr => serializeDataType(attr.dataType).isDefined)) { + withFallbackReason(scan, unsupportedDefaultReason) + } else { + // There are unsupported scan type + withFallbackReason( + scan, + s"unsupported Comet operator: ${scan.nodeName}, due to unsupported data types above") + } + None + } + } + + /** + * Build the `NativeScanCommon` proto shared by the core parquet scan and contrib scans that + * delegate to the same native parquet machinery (e.g. a Delta scan contrib, which passes + * physical-name schemas under column mapping). Returns `None` when an output data type or an + * existence default value cannot be serialized; the caller is responsible for tagging a + * fallback reason. + * + * Visibility note: `private[comet]` means a contrib caller must live under an + * `org.apache.comet.*` package (the same constraint `PlanDataInjector` implementers have). + */ + private[comet] def buildNativeScanCommon( + source: String, + output: Seq[Attribute], + requiredSchema: StructType, + dataSchema: StructType, + partitionSchema: StructType, + fileConstantMetadataColumns: Seq[AttributeReference], + dataFilters: Seq[Expression], + objectStoreOptions: Map[String, String], + hadoopConf: Configuration, + conf: SQLConf): Option[OperatorOuterClass.NativeScanCommon.Builder] = { val commonBuilder = OperatorOuterClass.NativeScanCommon.newBuilder() // Set source in common (used as part of injection key) - commonBuilder.setSource(scan.simpleStringWithNodeId()) + commonBuilder.setSource(source) - val scanTypes = scan.output.flatten { attr => + val scanTypes = output.flatten { attr => serializeDataType(attr.dataType) } - if (scanTypes.length == scan.output.length) { - // Collect S3/cloud storage configurations - val hadoopConf = scan.relation.sparkSession.sessionState - .newHadoopConfWithOptions(scan.relation.options) - - // Use selectedPartitions (static) instead of getFilePartitions() because at planning time - // DPP subqueries haven't been resolved yet. Object store options don't depend on DPP. - val fileUris = scan.selectedPartitions.view.flatMap(_.files).map(_.getPath.toUri) - val objectStoreOptions = - NativeConfig.extractScanObjectStoreOptions(hadoopConf, fileUris) match { - case Right(options) => options - case Left(reason) => - // CometExecRule falls back to the wrapped Spark scan, so it carries the reason too. - withFallbackReason(scan, reason) - withFallbackReason(scan.wrapped, reason) - return None + if (scanTypes.length != output.length) { + // There are unsupported scan types + return None + } + commonBuilder.addAllFields(scanTypes.asJava) + + if (conf.getConf(SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED)) { + commonBuilder.setHasDataFilters(dataFilters.nonEmpty) + val filterProtos = new ListBuffer[Expr]() + for (filter <- dataFilters) { + exprToProto(filter, output) match { + case Some(proto) => filterProtos += proto + case _ => + logWarning(s"Unsupported data filter $filter") } + } + commonBuilder.addAllDataFilters(filterProtos.asJava) + } + + serializeExistenceDefaultValues(requiredSchema, output) match { + case Some((defaultValues, indexes)) => + commonBuilder.addAllDefaultValues(defaultValues.asJava) + commonBuilder.addAllDefaultValuesIndexes(indexes.asJava) + case None => + // An unserializable existence default: fail closed rather than misalign the lists. + return None + } - commonBuilder.addAllFields(scanTypes.asJava) + // Constant metadata columns (file_path, file_name, file_size, file_block_start, + // file_block_length, file_modification_time) are known before opening the file and + // constant for every row read from it, exactly like partition columns. Spark places + // them immediately after partition columns in the scan output + // (FileSourceStrategy.scala: readDataColumns ++ generatedMetadataColumns ++ + // partitionColumns ++ constantMetadataColumns), so appending them after the real + // partition schema here keeps the two in lockstep. + val constantMetadataFields = uniqueConstantMetadataFields( + fileConstantMetadataColumns, + dataSchema.fields.map(_.name).toSet ++ partitionSchema.fields.map(_.name).toSet) + val partitionSchemaFields = partitionSchema.fields.toSeq ++ constantMetadataFields + val partitionSchemaProto = schema2Proto(partitionSchemaFields) + val requiredSchemaProto = schema2Proto(requiredSchema) + + // Serialize the same pruned schema `isSupported` validated; see `nativeDataSchema`. + val prunedDataSchema = nativeDataSchema(dataSchema, requiredSchema, conf.resolver) + val dataSchemaProto = schema2Proto(prunedDataSchema) + + val dataSchemaIndexes = requiredSchema.map(field => { + prunedDataSchema.fieldIndex(field.name) + }) + val partitionSchemaIndexes = prunedDataSchema.fields.length until + (prunedDataSchema.length + partitionSchemaFields.length) - // Sink operators don't have children - builder.clearChildren() + val projectionVector = (dataSchemaIndexes ++ partitionSchemaIndexes).map(idx => + idx.toLong.asInstanceOf[java.lang.Long]) - if (scan.conf.getConf(SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED)) { - val supportedDataFilters = scan.supportedDataFilters - commonBuilder.setHasDataFilters(supportedDataFilters.nonEmpty) - val dataFilters = new ListBuffer[Expr]() - for (filter <- supportedDataFilters) { - exprToProto(filter, scan.output) match { - case Some(proto) => dataFilters += proto - case _ => - logWarning(s"Unsupported data filter $filter") - } - } - commonBuilder.addAllDataFilters(dataFilters.asJava) - } + commonBuilder.addAllProjectionVector(projectionVector.asJava) - serializeExistenceDefaultValues(scan.requiredSchema, scan.output) match { - case Some((defaultValues, indexes)) => - commonBuilder.addAllDefaultValues(defaultValues.asJava) - commonBuilder.addAllDefaultValuesIndexes(indexes.asJava) - case None => - withFallbackReason(scan, unsupportedDefaultReason) - return None - } + // In `CometScanRule`, we ensure partitionSchema (including constant metadata columns) + // is supported. + assert(partitionSchemaProto.length == partitionSchemaFields.length) - // Constant metadata columns (file_path, file_name, file_size, file_block_start, - // file_block_length, file_modification_time) are known before opening the file and - // constant for every row read from it, exactly like partition columns. Spark places - // them immediately after partition columns in `scan.output` - // (FileSourceStrategy.scala: readDataColumns ++ generatedMetadataColumns ++ - // partitionColumns ++ constantMetadataColumns), so appending them after the real - // partition schema here keeps the two in lockstep. - val constantMetadataFields = uniqueConstantMetadataFields( - scan.wrapped.fileConstantMetadataColumns, - scan.relation.dataSchema.fields.map(_.name).toSet ++ - scan.relation.partitionSchema.fields.map(_.name).toSet) - val partitionSchemaFields = scan.relation.partitionSchema.fields.toSeq ++ - constantMetadataFields - val partitionSchema = schema2Proto(partitionSchemaFields) - val requiredSchema = schema2Proto(scan.requiredSchema) - - // Serialize the same pruned schema `isSupported` validated; see `nativeDataSchema`. - val prunedDataSchema = - nativeDataSchema(scan.relation.dataSchema, scan.requiredSchema, scan.conf.resolver) - val dataSchema = schema2Proto(prunedDataSchema) - - val dataSchemaIndexes = scan.requiredSchema.map(field => { - prunedDataSchema.fieldIndex(field.name) - }) - val partitionSchemaIndexes = prunedDataSchema.fields.length until - (prunedDataSchema.length + partitionSchemaFields.length) - - val projectionVector = (dataSchemaIndexes ++ partitionSchemaIndexes).map(idx => - idx.toLong.asInstanceOf[java.lang.Long]) - - commonBuilder.addAllProjectionVector(projectionVector.asJava) - - // In `CometScanRule`, we ensure partitionSchema (including constant metadata columns) - // is supported. - assert(partitionSchema.length == partitionSchemaFields.length) - - commonBuilder.addAllDataSchema(dataSchema.asJava) - commonBuilder.addAllRequiredSchema(requiredSchema.asJava) - commonBuilder.addAllPartitionSchema(partitionSchema.asJava) - val sessionTimeZone = scan.conf.getConfString("spark.sql.session.timeZone") - commonBuilder.setSessionTimezone( - CometTimeZone.nativeId(Some(sessionTimeZone)).getOrElse(sessionTimeZone)) - commonBuilder.setCaseSensitive(scan.conf.getConf[Boolean](SQLConf.CASE_SENSITIVE)) - - // SPARK-53535 (Spark 4.1+): when reading a struct whose requested fields are all - // missing in the Parquet file, the new default preserves the parent struct's - // nullness from the file (so non-null parents materialize as a struct of all-null - // fields). Pre-4.1 Spark hardcodes the legacy behavior (whole struct null), which - // matches the Comet default we use as fallback. - val returnNullStructConfKey = - "spark.sql.legacy.parquet.returnNullStructIfAllFieldsMissing" - val returnNullStructDefault = if (isSpark41Plus) "false" else "true" - commonBuilder.setReturnNullStructIfAllFieldsMissing( - scan.conf.getConfString(returnNullStructConfKey, returnNullStructDefault).toBoolean) - - // Field-ID matching: only ask the native side to do extra work when the conf is on AND - // the requested schema actually carries IDs. Spark's ParquetReadSupport applies the same - // gate before invoking matchIdField. - val hasFieldIds = ParquetUtils.hasFieldIds(scan.requiredSchema) - val useFieldId = scan.conf.getConf(SQLConf.PARQUET_FIELD_ID_READ_ENABLED) && hasFieldIds - commonBuilder.setUseFieldId(useFieldId) - // Spark's ParquetReadSupport refuses a file without field ids whenever the requested - // schema carries one, whatever the read flag says, unless ignoreMissing is set. - commonBuilder.setRequireFieldIds( - hasFieldIds && !scan.conf.getConf(SQLConf.IGNORE_MISSING_PARQUET_FIELD_ID)) - - commonBuilder.setAllowTypePromotion(CometConf.COMET_SCHEMA_EVOLUTION_ENABLED) - commonBuilder.setAllowTimestampLtzToNtz(CometConf.COMET_ALLOW_TIMESTAMP_LTZ_AS_NTZ) - - commonBuilder.setEncryptionEnabled(CometParquetUtils.encryptionEnabled(hadoopConf)) - - objectStoreOptions.foreach { case (key, value) => - commonBuilder.putObjectStoreOptions(key, value) - } + commonBuilder.addAllDataSchema(dataSchemaProto.asJava) + commonBuilder.addAllRequiredSchema(requiredSchemaProto.asJava) + commonBuilder.addAllPartitionSchema(partitionSchemaProto.asJava) - // Set common data in NativeScan (file_partition will be populated at execution time) - nativeScanBuilder.setCommon(commonBuilder.build()) + populateScanConfFlags(commonBuilder, requiredSchema, objectStoreOptions, hadoopConf, conf) - Some(builder.setNativeScan(nativeScanBuilder).build()) + Some(commonBuilder) + } - } else { - // There are unsupported scan type - withFallbackReason( - scan, - s"unsupported Comet operator: ${scan.nodeName}, due to unsupported data types above") - None + /** + * Populate the configuration-derived flags of a `NativeScanCommon`: session timezone, case + * sensitivity, struct-nullness legacy flag, field-ID matching, type promotion, encryption, and + * object-store options. Shared with contrib scans that assemble their own schemas/projection + * (e.g. the Delta contrib's deletion-vector shape) so new flags added here reach them without + * drift. + */ + private[comet] def populateScanConfFlags( + commonBuilder: OperatorOuterClass.NativeScanCommon.Builder, + requiredSchema: StructType, + objectStoreOptions: Map[String, String], + hadoopConf: Configuration, + conf: SQLConf): Unit = { + val sessionTimeZone = conf.getConfString("spark.sql.session.timeZone") + commonBuilder.setSessionTimezone( + CometTimeZone.nativeId(Some(sessionTimeZone)).getOrElse(sessionTimeZone)) + commonBuilder.setCaseSensitive(conf.getConf[Boolean](SQLConf.CASE_SENSITIVE)) + + // SPARK-53535 (Spark 4.1+): when reading a struct whose requested fields are all + // missing in the Parquet file, the new default preserves the parent struct's + // nullness from the file (so non-null parents materialize as a struct of all-null + // fields). Pre-4.1 Spark hardcodes the legacy behavior (whole struct null), which + // matches the Comet default we use as fallback. + val returnNullStructConfKey = + "spark.sql.legacy.parquet.returnNullStructIfAllFieldsMissing" + val returnNullStructDefault = if (isSpark41Plus) "false" else "true" + commonBuilder.setReturnNullStructIfAllFieldsMissing( + conf.getConfString(returnNullStructConfKey, returnNullStructDefault).toBoolean) + + // Field-ID matching: only ask the native side to do extra work when the conf is on AND + // the requested schema actually carries IDs. Spark's ParquetReadSupport applies the same + // gate before invoking matchIdField. + val hasFieldIds = ParquetUtils.hasFieldIds(requiredSchema) + val useFieldId = conf.getConf(SQLConf.PARQUET_FIELD_ID_READ_ENABLED) && hasFieldIds + commonBuilder.setUseFieldId(useFieldId) + // Spark's ParquetReadSupport refuses a file without field ids whenever the requested + // schema carries one, whatever the read flag says, unless ignoreMissing is set. + commonBuilder.setRequireFieldIds( + hasFieldIds && !conf.getConf(SQLConf.IGNORE_MISSING_PARQUET_FIELD_ID)) + + commonBuilder.setAllowTypePromotion(CometConf.COMET_SCHEMA_EVOLUTION_ENABLED) + commonBuilder.setAllowTimestampLtzToNtz(CometConf.COMET_ALLOW_TIMESTAMP_LTZ_AS_NTZ) + + commonBuilder.setEncryptionEnabled(CometParquetUtils.encryptionEnabled(hadoopConf)) + + objectStoreOptions.foreach { case (key, value) => + commonBuilder.putObjectStoreOptions(key, value) } - } override def createExec(nativeOp: Operator, op: CometScanExec): CometNativeExec = { CometNativeScanExec(nativeOp, op.wrapped, op.session, op) } + + /** + * Sets the `inline_data` bytes field on a `DeltaSparkDvDescriptor` builder. The shade plugin + * relocates `com.google.protobuf.ByteString` when packaged, rewriting bytecode descriptors but + * not a Scala method's own pickled signature, so a helper returning `ByteString` directly would + * disagree with the packaged jar's Java-generated `setInlineData(ByteString)`. Keeping the + * protobuf type out of this method's signature sidesteps that, letting out-of-tree modules + * (e.g. Delta contrib) call this whether compiled against unshaded or shaded classes. + */ + def setDvInlineData( + builder: OperatorOuterClass.DeltaSparkDvDescriptor.Builder, + bytes: Array[Byte]): OperatorOuterClass.DeltaSparkDvDescriptor.Builder = + builder.setInlineData(com.google.protobuf.ByteString.copyFrom(bytes)) } diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/package.scala b/spark/src/main/scala/org/apache/comet/serde/operator/package.scala index cf6e3fabe8d..bee7f61a128 100644 --- a/spark/src/main/scala/org/apache/comet/serde/operator/package.scala +++ b/spark/src/main/scala/org/apache/comet/serde/operator/package.scala @@ -106,7 +106,9 @@ package object operator { // In `CometScanRule`, we have already checked that all partition and metadata column values // are supported. So, we can safely use `get` here. - private def literalToProto(literal: Literal, description: String): ExprOuterClass.Expr = { + private[comet] def literalToProto( + literal: Literal, + description: String): ExprOuterClass.Expr = { val valueProto = exprToProto(literal, Seq.empty) assert(valueProto.isDefined, s"Unsupported $description") valueProto.get diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala index db1ca44c200..d20022f2b53 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala @@ -69,7 +69,11 @@ private[spark] class CometExecRDD( encryptedFilePaths: Seq[String] = Seq.empty, shuffleScanIndices: Set[Int] = Set.empty, @transient perPartitionFilePaths: Array[Seq[String]] = Array.empty, - outputTypes: Seq[DataType] = Seq.empty) + outputTypes: Seq[DataType] = Seq.empty, + // Set by a contrib leaf scan (e.g. the Delta contrib's `CometDeltaNativeScanExec`) that + // builds this RDD directly, bypassing `CometNativeExec.executeColumnarWithContext`'s own + // `ctx.hasScanInput` check, so it reports task input metrics without subclassing this RDD. + reportScanInputMetrics: Boolean = false) extends RDD[ColumnarBatch](sc, inputRDDs.map(rdd => new OneToOneDependency(rdd))) { // Determine partition count: from inputs if available, otherwise from parameter @@ -104,6 +108,11 @@ private[spark] class CometExecRDD( // reverse registration order, so registering first means this listener runs last, after // nested native blocks and the iterator have published their final metric values. Option(context).foreach(nativeMetrics.reportSpillMetrics) + // Registered here for the same reason: it has to run after the iterator's close has + // published the final scan metrics. + if (reportScanInputMetrics) { + Option(context).foreach(nativeMetrics.reportScanInputMetrics) + } val partition = split.asInstanceOf[CometExecPartition] @@ -233,7 +242,8 @@ object CometExecRDD { encryptedFilePaths: Seq[String] = Seq.empty, shuffleScanIndices: Set[Int] = Set.empty, perPartitionFilePaths: Array[Seq[String]] = Array.empty, - outputTypes: Seq[DataType] = Seq.empty): CometExecRDD = { + outputTypes: Seq[DataType] = Seq.empty, + reportScanInputMetrics: Boolean = false): CometExecRDD = { // scalastyle:on new CometExecRDD( @@ -251,6 +261,7 @@ object CometExecRDD { encryptedFilePaths, shuffleScanIndices, perPartitionFilePaths, - outputTypes) + outputTypes, + reportScanInputMetrics) } } 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 de354502c59..a1cc6f04c41 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 @@ -1179,9 +1179,10 @@ abstract class CometNativeExec extends CometExec { commonByKey = commonByKey, perPartitionByKey = perPartitionByKey, shuffleScanIndices = shuffleScanIndices, - // A leaf Comet scan (`CometNativeScanExec`, `CometIcebergNativeScanExec`) can - // contribute `bytes_scanned` / `output_rows` to Spark's task-level input metrics, - // which drive the Input column on the UI's Stages and Executors tabs. + // A leaf Comet scan (`CometNativeScanExec`, `CometIcebergNativeScanExec`, or a contrib + // leaf such as `CometDeltaNativeScanExec`) can contribute `bytes_scanned` / + // `output_rows` to Spark's task-level input metrics, which drive the Input column on + // the UI's Stages and Executors tabs. // Matching on `CometLeafExec` rather than `CometNativeScanExec` keeps every scan // reported once the scan is fused into a larger native block, where only the block // root's `compute` runs. `reportScanInputMetrics` self-filters on the `bytes_scanned` diff --git a/spark/src/test/scala/org/apache/comet/CometS3TestBase.scala b/spark/src/test/scala/org/apache/comet/CometS3TestBase.scala index a15e3f75be0..07186b03550 100644 --- a/spark/src/test/scala/org/apache/comet/CometS3TestBase.scala +++ b/spark/src/test/scala/org/apache/comet/CometS3TestBase.scala @@ -75,6 +75,9 @@ trait CometS3TestBase extends CometTestBase { conf.set("spark.hadoop.fs.s3a.secret.key", password) conf.set("spark.hadoop.fs.s3a.endpoint", minioContainer.getS3URL) conf.set("spark.hadoop.fs.s3a.path.style.access", "true") + // Pin the region explicitly rather than relying on Hadoop-version-dependent region + // resolution; MinIO ignores the value. Native maps this the same way (see s3.rs). + conf.set("spark.hadoop.fs.s3a.endpoint.region", "us-east-1") } // Spark 4.2 has no published Iceberg spark-runtime yet; the build reuses the 4.0 runtime, whose diff --git a/spark/src/test/scala/org/apache/comet/rules/CometScanContribSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometScanContribSuite.scala index c848501a71a..4e41a9b147a 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometScanContribSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometScanContribSuite.scala @@ -26,10 +26,15 @@ import java.nio.file.Files import java.util.ServiceLoader import scala.annotation.nowarn +import scala.collection.mutable.ArrayBuffer import scala.jdk.CollectionConverters._ import org.scalatest.funsuite.AnyFunSuite +import org.apache.logging.log4j.LogManager +import org.apache.logging.log4j.core.LogEvent +import org.apache.logging.log4j.core.appender.AbstractAppender +import org.apache.logging.log4j.core.config.Property import org.apache.spark.rdd.RDD import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.InternalRow @@ -38,6 +43,7 @@ import org.apache.spark.sql.execution.{FileSourceScanExec, LeafExecNode, SparkPl import org.apache.spark.sql.execution.datasources.HadoopFsRelation import org.apache.spark.sql.execution.datasources.v2.BatchScanExec +import org.apache.comet.ContribServices import org.apache.comet.util.ClassLoaders /** @@ -108,6 +114,35 @@ class CometScanContribSuite extends AnyFunSuite { assert(CometScanContrib.tryTransformV2(null).isEmpty) } + test("a provider whose class cannot link is skipped by discovery and the rest still load") { + // ServiceLoader raises NoClassDefFoundError straight out of Class.forName when a listed + // provider's superclass or interface is missing, the version-skewed-jar case. Discovery + // must log and skip it so the remaining providers are found and no scan throws. + val skewed = "org.apache.comet.rules.SkewedProvider" + withServiceFile(Seq(skewed, classOf[ClaimingScanContrib].getName)) { fileLoader => + val loader = new ClassLoader(fileLoader) { + override def loadClass(name: String, resolve: Boolean): Class[_] = { + if (name == skewed) { + throw new NoClassDefFoundError("org/apache/comet/rules/MissingContribInterface") + } + super.loadClass(name, resolve) + } + } + val events = withCapturedLogEvents(ContribServices.getClass.getName.stripSuffix("$")) { + val discovered = CometScanContrib.loadContribs(loader) + assert( + discovered.exists(_.isInstanceOf[ClaimingScanContrib]), + "the linkable provider should still be discovered, got: " + + discovered.map(_.getClass.getName)) + assert(!discovered.exists(_.getClass.getName == skewed)) + } + val messages = events.map(_.getMessage.getFormattedMessage) + assert( + messages.exists(m => m.contains(classOf[NoClassDefFoundError].getName)), + s"expected a warning naming the LinkageError subtype, got: $messages") + } + } + test("a contrib registered via META-INF/services is discovered and its claim is returned") { // Proves the whole registration contract a contrib depends on: dropping a service file naming // an implementation makes it visible to ServiceLoader, and a Some(...) it returns is what the @@ -245,10 +280,63 @@ class CometScanContribSuite extends AnyFunSuite { } test("a fatal error from a contrib is not swallowed") { - // NonFatal deliberately lets LinkageError/OOM-class failures through: those signal a broken - // JVM or a mis-built jar, not a scan this contrib cannot plan. + // OutOfMemoryError is neither NonFatal nor a LinkageError: it signals real JVM-level + // exhaustion, not a version-skewed contrib jar, and must always propagate uncontained. val contribs = Seq(new FatalScanContrib) - intercept[LinkageError](offerV1(contribs)) + intercept[OutOfMemoryError](offerV1(contribs)) + } + + test( + "a LinkageError from a contrib is contained, logged by name, and the next contrib still " + + "gets a look") { + // A version-skewed contrib jar (compiled against a Comet internal that has since moved or + // been removed) throws NoSuchMethodError/NoClassDefFoundError -- a LinkageError, which + // NonFatal does not match. It must be contained the same way a NonFatal decline is: logged, + // treated as "does not claim this scan", and the next contrib still consulted. + val contribs = Seq(new VersionSkewedScanContrib, new ClaimingScanContrib) + val events = withCapturedLogEvents(classOf[CometScanContrib].getName) { + assert(offerV1(contribs).contains(ContribStubs.ClaimedByV1)) + assert(offerV2(contribs).contains(ContribStubs.ClaimedByV2)) + } + val messages = events.map(_.getMessage.getFormattedMessage) + assert( + messages.count(m => + m.contains(classOf[VersionSkewedScanContrib].getName) && + m.contains(classOf[NoSuchMethodError].getName)) == 2, + "expected one warning per hook naming both the contrib class and the LinkageError " + + s"subtype, got: $messages") + } + + test("a LinkageError with nothing behind it declines rather than failing the query") { + val contribs = Seq(new VersionSkewedScanContrib) + assert(offerV1(contribs).isEmpty, "the scan must fall through to Comet's built-in handling") + assert(offerV2(contribs).isEmpty) + } + + /** + * Attaches a minimal Log4j2 appender directly to the logger named `loggerName` for the duration + * of `f`, returning every event it captured. `CometScanContrib`'s `logWarning` calls go through + * Spark's `Logging` trait to a logger named after the emitting class, so this lets a test + * assert a specific warning was actually emitted -- not merely that the surrounding code path + * didn't throw. Restores the logger's prior appenders/level afterward so this cannot leak into + * other tests in the same JVM. + */ + private def withCapturedLogEvents(loggerName: String)(f: => Unit): Seq[LogEvent] = { + val logger = + LogManager.getLogger(loggerName).asInstanceOf[org.apache.logging.log4j.core.Logger] + val appender = new CapturingAppender(s"CometScanContribSuite-${System.nanoTime()}") + appender.start() + val originalLevel = logger.getLevel + logger.addAppender(appender) + logger.setLevel(org.apache.logging.log4j.Level.WARN) + try { + f + appender.events.toSeq + } finally { + logger.removeAppender(appender) + logger.setLevel(originalLevel) + appender.stop() + } } /** @@ -356,12 +444,45 @@ class ThrowingScanContrib extends CometScanContrib { throw new IllegalStateException("contrib blew up while planning a V2 scan") } -/** Fails in a way that must NOT be caught. */ +/** Fails in a way that must NOT be caught: neither `NonFatal` nor a `LinkageError`. */ class FatalScanContrib extends CometScanContrib { override def tryTransformV1( plan: SparkPlan, session: SparkSession, scanExec: FileSourceScanExec, relation: HadoopFsRelation): Option[SparkPlan] = - throw new NoClassDefFoundError("mis-built contrib jar") + throw new OutOfMemoryError("simulated JVM-level exhaustion, not a version-skewed contrib jar") +} + +/** + * Simulates a contrib jar built against a Comet internal (a method signature, a class) that has + * since moved, been renamed, or been removed -- the exact failure mode a stale `--jars` contrib + * hits against a newer Comet on the driver's classpath. Must be contained the same way a + * `NonFatal` decline is, unlike [[FatalScanContrib]]'s genuinely fatal error. + */ +class VersionSkewedScanContrib extends CometScanContrib { + override def tryTransformV1( + plan: SparkPlan, + session: SparkSession, + scanExec: FileSourceScanExec, + relation: HadoopFsRelation): Option[SparkPlan] = + throw new NoSuchMethodError( + "org.apache.comet.rules.CometScanContribSuite$InternalApi.movedMethod()V") + + override def tryTransformV2(scanExec: BatchScanExec): Option[SparkPlan] = + throw new NoSuchMethodError( + "org.apache.comet.rules.CometScanContribSuite$InternalApi.movedMethod()V") +} + +/** + * Minimal Log4j2 appender that records every event it receives, verbatim, for + * [[CometScanContribSuite.withCapturedLogEvents]] to inspect after the fact. + */ +private class CapturingAppender(name: String) + extends AbstractAppender(name, null, null, false, Property.EMPTY_ARRAY) { + val events: ArrayBuffer[LogEvent] = ArrayBuffer.empty + + override def append(event: LogEvent): Unit = events.synchronized { + events += event.toImmutable + } }